refactor: reorganize the sqlglot scalar compiler layout - part 2 (#2093)

This change follows up on #2075 by splitting unary_compiler.py and its unit test file into smaller files.

Author: chelsea-lin

bigframes/core/compile/sqlglot/__init__.py

@@ -14,7 +14,18 @@
 from __future__ import annotations
 
 from bigframes.core.compile.sqlglot.compiler import SQLGlotCompiler
+import bigframes.core.compile.sqlglot.expressions.array_ops # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.binary_compiler # noqa: F401
-import bigframes.core.compile.sqlglot.expressions.unary_compiler # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.blob_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.comparison_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.date_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.datetime_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.generic_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.geo_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.json_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.numeric_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.string_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.struct_ops # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.timedelta_ops # noqa: F401
 
 __all__ = ["SQLGlotCompiler"]

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

@@ -0,0 +1,68 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import typing
+
+import sqlglot
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+
+register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(ops.ArrayToStringOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArrayToStringOp) -> sge.Expression:
+ return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
+
+
+@register_unary_op(ops.ArrayIndexOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArrayIndexOp) -> sge.Expression:
+ return sge.Bracket(
+ this=expr.expr,
+ expressions=[sge.Literal.number(op.index)],
+ safe=True,
+ offset=False,
+ )
+
+
+@register_unary_op(ops.ArraySliceOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
+ slice_idx = sqlglot.to_identifier("slice_idx")
+
+ conditions: typing.List[sge.Predicate] = [slice_idx >= op.start]
+
+ if op.stop is not None:
+ conditions.append(slice_idx < op.stop)
+
+ # local name for each element in the array
+ el = sqlglot.to_identifier("el")
+
+ selected_elements = (
+ sge.select(el)
+ .from_(
+ sge.Unnest(
+ expressions=[expr.expr],
+ alias=sge.TableAlias(columns=[el]),
+ offset=slice_idx,
+ )
+ )
+ .where(*conditions)
+ )
+
+ return sge.array(selected_elements)

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

@@ -0,0 +1,33 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+
+register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(ops.obj_fetch_metadata_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("OBJ.FETCH_METADATA", expr.expr)
+
+
+@register_unary_op(ops.ObjGetAccessUrl)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("OBJ.GET_ACCESS_URL", expr.expr)

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

@@ -0,0 +1,59 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import typing
+
+import pandas as pd
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+import bigframes.dtypes as dtypes
+
+register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(ops.IsInOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
+ values = []
+ is_numeric_expr = dtypes.is_numeric(expr.dtype)
+ for value in op.values:
+ if value is None:
+ continue
+ dtype = dtypes.bigframes_type(type(value))
+ if expr.dtype == dtype or is_numeric_expr and dtypes.is_numeric(dtype):
+ values.append(sge.convert(value))
+
+ if op.match_nulls:
+ contains_nulls = any(_is_null(value) for value in op.values)
+ if contains_nulls:
+ return sge.Is(this=expr.expr, expression=sge.Null()) | sge.In(
+ this=expr.expr, expressions=values
+ )
+
+ if len(values) == 0:
+ return sge.convert(False)
+
+ return sge.func(
+ "COALESCE", sge.In(this=expr.expr, expressions=values), sge.convert(False)
+ )
+
+
+# Helpers
+def _is_null(value) -> bool:
+ # float NaN/inf should be treated as distinct from 'true' null values
+ return typing.cast(bool, pd.isna(value)) and not isinstance(value, float)

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

@@ -0,0 +1,61 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+
+register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(ops.date_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Date(this=expr.expr)
+
+
+@register_unary_op(ops.day_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="DAY"), expression=expr.expr)
+
+
+@register_unary_op(ops.dayofweek_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ # Adjust the 1-based day-of-week index (from SQL) to a 0-based index.
+ return sge.Extract(
+ this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr
+ ) - sge.convert(1)
+
+
+@register_unary_op(ops.dayofyear_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="DAYOFYEAR"), expression=expr.expr)
+
+
+@register_unary_op(ops.iso_day_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr)
+
+
+@register_unary_op(ops.iso_week_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="ISOWEEK"), expression=expr.expr)
+
+
+@register_unary_op(ops.iso_year_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="ISOYEAR"), expression=expr.expr)

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

@@ -0,0 +1,99 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+
+register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(ops.FloorDtOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.FloorDtOp) -> 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))
+
+
+@register_unary_op(ops.hour_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="HOUR"), expression=expr.expr)
+
+
+@register_unary_op(ops.minute_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="MINUTE"), expression=expr.expr)
+
+
+@register_unary_op(ops.month_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="MONTH"), expression=expr.expr)
+
+
+@register_unary_op(ops.normalize_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this="DAY"))
+
+
+@register_unary_op(ops.quarter_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="QUARTER"), expression=expr.expr)
+
+
+@register_unary_op(ops.second_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="SECOND"), expression=expr.expr)
+
+
+@register_unary_op(ops.StrftimeOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.StrftimeOp) -> sge.Expression:
+ return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr)
+
+
+@register_unary_op(ops.time_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("TIME", expr.expr)
+
+
+@register_unary_op(ops.ToDatetimeOp)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Cast(this=sge.func("TIMESTAMP_SECONDS", expr.expr), to="DATETIME")
+
+
+@register_unary_op(ops.ToTimestampOp)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("TIMESTAMP_SECONDS", expr.expr)
+
+
+@register_unary_op(ops.UnixMicros)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("UNIX_MICROS", expr.expr)
+
+
+@register_unary_op(ops.UnixMillis)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.func("UNIX_MILLIS", expr.expr)
+
+
+@register_unary_op(ops.UnixSeconds, pass_op=True)
+def _(expr: TypedExpr, op: ops.UnixSeconds) -> sge.Expression:
+ return sge.func("UNIX_SECONDS", expr.expr)
+
+
+@register_unary_op(ops.year_op)
+def _(expr: TypedExpr) -> sge.Expression:
+ return sge.Extract(this=sge.Identifier(this="YEAR"), expression=expr.expr)