Merge branch 'main' into main_chelsealin_add

Author: chelsea-lin

bigframes/core/compile/sqlglot/expressions/unary_compiler.py

@@ -16,6 +16,8 @@
 
 import typing
 
+import pandas as pd
+import pyarrow as pa
 import sqlglot
 import sqlglot.expressions as sge
 
@@ -105,6 +107,12 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     )
 
 
+@UNARY_OP_REGISTRATION.register(ops.AsTypeOp)
+def _(op: ops.AsTypeOp, expr: TypedExpr) -> sge.Expression:
+    # TODO: Support more types for casting, such as JSON, etc.
+    return sge.Cast(this=expr.expr, to=op.to_type)
+
+
 @UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
 def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
     return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
@@ -234,6 +242,12 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     ) - sge.convert(1)
 
 
+@UNARY_OP_REGISTRATION.register(ops.FloorDtOp)
+def _(op: ops.FloorDtOp, expr: TypedExpr) -> sge.Expression:
+    # TODO: Remove this method when it is covered by ops.FloorOp
+    return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this=op.freq))
+
+
 @UNARY_OP_REGISTRATION.register(ops.floor_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.Floor(this=expr.expr)
@@ -249,6 +263,26 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("ST_ASTEXT", expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.geo_st_boundary_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_BOUNDARY", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.geo_st_geogfromtext_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("SAFE.ST_GEOGFROMTEXT", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.geo_st_isclosed_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_ISCLOSED", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.GeoStLengthOp)
+def _(op: ops.GeoStLengthOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_LENGTH", expr.expr)
+
+
 @UNARY_OP_REGISTRATION.register(ops.geo_x_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("SAFE.ST_X", expr.expr)
@@ -274,6 +308,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.BitwiseNot(this=expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.IsInOp)
+def _(op: ops.IsInOp, expr: TypedExpr) -> sge.Expression:
+    return sge.In(this=expr.expr, expressions=[sge.convert(v) for v in op.values])
+
+
 @UNARY_OP_REGISTRATION.register(ops.isalnum_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.RegexpLike(this=expr.expr, expression=sge.convert(r"^(\p{N}|\p{L})+$"))
@@ -517,6 +556,26 @@ def _(op: ops.StrSliceOp, expr: TypedExpr) -> sge.Expression:
     )
 
 
+@UNARY_OP_REGISTRATION.register(ops.StrftimeOp)
+def _(op: ops.StrftimeOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.StructFieldOp)
+def _(op: ops.StructFieldOp, expr: TypedExpr) -> sge.Expression:
+    if isinstance(op.name_or_index, str):
+        name = op.name_or_index
+    else:
+        pa_type = typing.cast(pd.ArrowDtype, expr.dtype)
+        pa_struct_type = typing.cast(pa.StructType, pa_type.pyarrow_dtype)
+        name = pa_struct_type.field(op.name_or_index).name
+
+    return sge.Column(
+        this=sge.to_identifier(name, quoted=True),
+        catalog=expr.expr,
+    )
+
+
 @UNARY_OP_REGISTRATION.register(ops.tan_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("TAN", expr.expr)
@@ -537,6 +596,36 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.Floor(this=expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.ToDatetimeOp)
+def _(op: ops.ToDatetimeOp, expr: TypedExpr) -> sge.Expression:
+    return sge.Cast(this=sge.func("TIMESTAMP_SECONDS", expr.expr), to="DATETIME")
+
+
+@UNARY_OP_REGISTRATION.register(ops.ToTimestampOp)
+def _(op: ops.ToTimestampOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("TIMESTAMP_SECONDS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.ToTimedeltaOp)
+def _(op: ops.ToTimedeltaOp, expr: TypedExpr) -> sge.Expression:
+    return sge.Interval(this=expr.expr, unit=sge.Identifier(this="SECOND"))
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixMicros)
+def _(op: ops.UnixMicros, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_MICROS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixMillis)
+def _(op: ops.UnixMillis, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_MILLIS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixSeconds)
+def _(op: ops.UnixSeconds, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_SECONDS", expr.expr)
+
+
 # JSON Ops
 @UNARY_OP_REGISTRATION.register(ops.JSONExtract)
 def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:

bigframes/pandas/io/api.py

@@ -16,6 +16,7 @@
 
 import functools
 import inspect
+import os
 import threading
 import typing
 from typing import (
@@ -56,6 +57,7 @@
 from bigframes.session import dry_runs
 import bigframes.session._io.bigquery
 import bigframes.session.clients
+import bigframes.session.metrics
 
 # Note: the following methods are duplicated from Session. This duplication
 # enables the following:
@@ -625,6 +627,11 @@ def _get_bqclient() -> bigquery.Client:
 
 def _dry_run(query, bqclient) -> bigquery.QueryJob:
     job = bqclient.query(query, bigquery.QueryJobConfig(dry_run=True))
+
+    # Fix for b/435183833. Log metrics even if a Session isn't available.
+    if bigframes.session.metrics.LOGGING_NAME_ENV_VAR in os.environ:
+        metrics = bigframes.session.metrics.ExecutionMetrics()
+        metrics.count_job_stats(job)
     return job
 
 

bigframes/session/metrics.py

@@ -40,32 +40,54 @@ def count_job_stats(
     ):
         if query_job is None:
             assert row_iterator is not None
-            total_bytes_processed = getattr(row_iterator, "total_bytes_processed", None)
-            query = getattr(row_iterator, "query", None)
-            if total_bytes_processed is None or query is None:
-                return
+
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = getattr(row_iterator, "total_bytes_processed", 0)
+            query_char_count = len(getattr(row_iterator, "query", ""))
+            slot_millis = getattr(row_iterator, "slot_millis", 0)
+            exec_seconds = 0.0
 
             self.execution_count += 1
-            self.query_char_count += len(query)
-            self.bytes_processed += total_bytes_processed
-            write_stats_to_disk(len(query), total_bytes_processed)
-            return
+            self.query_char_count += query_char_count
+            self.bytes_processed += bytes_processed
+            self.slot_millis += slot_millis
+
+        elif query_job.configuration.dry_run:
+            query_char_count = len(query_job.query)
 
-        if query_job.configuration.dry_run:
-            write_stats_to_disk(len(query_job.query), 0, 0, 0)
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = 0
+            slot_millis = 0
+            exec_seconds = 0.0
 
-        stats = get_performance_stats(query_job)
-        if stats is not None:
-            query_char_count, bytes_processed, slot_millis, execution_secs = stats
+        elif (stats := get_performance_stats(query_job)) is not None:
+            query_char_count, bytes_processed, slot_millis, exec_seconds = stats
             self.execution_count += 1
             self.query_char_count += query_char_count
             self.bytes_processed += bytes_processed
             self.slot_millis += slot_millis
-            self.execution_secs += execution_secs
+            self.execution_secs += exec_seconds
             write_stats_to_disk(
-                query_char_count, bytes_processed, slot_millis, execution_secs
+                query_char_count=query_char_count,
+                bytes_processed=bytes_processed,
+                slot_millis=slot_millis,
+                exec_seconds=exec_seconds,
             )
 
+        else:
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = 0
+            query_char_count = 0
+            slot_millis = 0
+            exec_seconds = 0
+
+        write_stats_to_disk(
+            query_char_count=query_char_count,
+            bytes_processed=bytes_processed,
+            slot_millis=slot_millis,
+            exec_seconds=exec_seconds,
+        )
+
 
 def get_performance_stats(
     query_job: bigquery.QueryJob,
@@ -103,10 +125,11 @@ def get_performance_stats(
 
 
 def write_stats_to_disk(
+    *,
     query_char_count: int,
     bytes_processed: int,
-    slot_millis: Optional[int] = None,
-    exec_seconds: Optional[float] = None,
+    slot_millis: int,
+    exec_seconds: float,
 ):
     """For pytest runs only, log information about the query job
     to a file in order to create a performance report.
@@ -118,18 +141,17 @@ def write_stats_to_disk(
     test_name = os.environ[LOGGING_NAME_ENV_VAR]
     current_directory = os.getcwd()
 
-    if (slot_millis is not None) and (exec_seconds is not None):
-        # store slot milliseconds
-        slot_file = os.path.join(current_directory, test_name + ".slotmillis")
-        with open(slot_file, "a") as f:
-            f.write(str(slot_millis) + "\n")
+    # store slot milliseconds
+    slot_file = os.path.join(current_directory, test_name + ".slotmillis")
+    with open(slot_file, "a") as f:
+        f.write(str(slot_millis) + "\n")
 
-        # store execution time seconds
-        exec_time_file = os.path.join(
-            current_directory, test_name + ".bq_exec_time_seconds"
-        )
-        with open(exec_time_file, "a") as f:
-            f.write(str(exec_seconds) + "\n")
+    # store execution time seconds
+    exec_time_file = os.path.join(
+        current_directory, test_name + ".bq_exec_time_seconds"
+    )
+    with open(exec_time_file, "a") as f:
+        f.write(str(exec_seconds) + "\n")
 
     # store length of query
     query_char_count_file = os.path.join(

tests/benchmark/read_gbq_colab/aggregate_output.py

@@ -48,7 +48,7 @@ def aggregate_output(*, project_id, dataset_id, table_id):
 
 
 if __name__ == "__main__":
-    config = utils.get_configuration(include_table_id=True)
+    config = utils.get_configuration(include_table_id=True, start_session=False)
     current_path = pathlib.Path(__file__).absolute()
 
     utils.get_execution_time(
@@ -58,5 +58,4 @@ def aggregate_output(*, project_id, dataset_id, table_id):
         project_id=config.project_id,
         dataset_id=config.dataset_id,
         table_id=config.table_id,
-        session=config.session,
     )

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_floor_dt/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `timestamp_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    TIMESTAMP_TRUNC(`bfcol_0`, DAY) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `timestamp_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_boundary/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_BOUNDARY(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_geogfromtext/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    SAFE.ST_GEOGFROMTEXT(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `string_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_isclosed/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_ISCLOSED(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_length/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_LENGTH(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_is_in/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` IN (1, 2, 3) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `int64_col`
+FROM `bfcte_1`