Merge branch 'main' into shuowei-fix-compiler-syntax-guards

Author: shuoweil

bigframes/core/blocks.py

@@ -1822,9 +1822,9 @@ def melt(
         Arguments correspond to pandas.melt arguments.
         """
         # TODO: Implement col_level and ignore_index
-        value_labels: pd.Index = pd.Index(
-            [self.col_id_to_label[col_id] for col_id in value_vars]
-        )
+        value_labels: pd.Index = self.column_labels[
+            [self.value_columns.index(col_id) for col_id in value_vars]
+        ]
         id_labels = [self.col_id_to_label[col_id] for col_id in id_vars]
 
         unpivot_expr, (var_col_ids, unpivot_out, passthrough_cols) = unpivot(
@@ -3417,6 +3417,7 @@ def unpivot(
         joined_array, (labels_mapping, column_mapping) = labels_array.relational_join(
             array_value, type="cross"
         )
+
     new_passthrough_cols = [column_mapping[col] for col in passthrough_columns]
     # Last column is offsets
     index_col_ids = [labels_mapping[col] for col in labels_array.column_ids[:-1]]
@@ -3426,20 +3427,24 @@ def unpivot(
     unpivot_exprs: List[ex.Expression] = []
     # Supports producing multiple stacked ouput columns for stacking only part of hierarchical index
     for input_ids in unpivot_columns:
-        # row explode offset used to choose the input column
-        # we use offset instead of label as labels are not necessarily unique
-        cases = itertools.chain(
-            *(
-                (
-                    ops.eq_op.as_expr(explode_offsets_id, ex.const(i)),
-                    ex.deref(column_mapping[id_or_null])
-                    if (id_or_null is not None)
-                    else ex.const(None),
+        col_expr: ex.Expression
+        if not input_ids:
+            col_expr = ex.const(None, dtype=bigframes.dtypes.INT_DTYPE)
+        else:
+            # row explode offset used to choose the input column
+            # we use offset instead of label as labels are not necessarily unique
+            cases = itertools.chain(
+                *(
+                    (
+                        ops.eq_op.as_expr(explode_offsets_id, ex.const(i)),
+                        ex.deref(column_mapping[id_or_null])
+                        if (id_or_null is not None)
+                        else ex.const(None),
+                    )
+                    for i, id_or_null in enumerate(input_ids)
                 )
-                for i, id_or_null in enumerate(input_ids)
             )
-        )
-        col_expr = ops.case_when_op.as_expr(*cases)
+            col_expr = ops.case_when_op.as_expr(*cases)
         unpivot_exprs.append(col_expr)
 
     joined_array, unpivot_col_ids = joined_array.compute_values(unpivot_exprs)
@@ -3457,19 +3462,43 @@ def _pd_index_to_array_value(
     Create an ArrayValue from a list of label tuples.
     The last column will be row offsets.
     """
+    id_gen = bigframes.core.identifiers.standard_id_strings()
+    col_ids = [next(id_gen) for _ in range(index.nlevels)]
+    offset_id = next(id_gen)
+
     rows = []
     labels_as_tuples = utils.index_as_tuples(index)
     for row_offset in range(len(index)):
-        id_gen = bigframes.core.identifiers.standard_id_strings()
         row_label = labels_as_tuples[row_offset]
         row_label = (row_label,) if not isinstance(row_label, tuple) else row_label
         row = {}
-        for label_part, id in zip(row_label, id_gen):
-            row[id] = label_part if pd.notnull(label_part) else None
-        row[next(id_gen)] = row_offset
+        for label_part, col_id in zip(row_label, col_ids):
+            row[col_id] = label_part if pd.notnull(label_part) else None
+        row[offset_id] = row_offset
         rows.append(row)
 
-    return core.ArrayValue.from_pyarrow(pa.Table.from_pylist(rows), session=session)
+    if not rows:
+        dtypes_list = getattr(index, "dtypes", None)
+        if dtypes_list is None:
+            dtypes_list = (
+                [index.dtype] if hasattr(index, "dtype") else [pd.Float64Dtype()]
+            )
+
+        fields = []
+        for col_id, dtype in zip(col_ids, dtypes_list):
+            try:
+                pa_type = bigframes.dtypes.bigframes_dtype_to_arrow_dtype(dtype)
+            except Exception:
+                pa_type = pa.string()
+            fields.append(pa.field(col_id, pa_type))
+        fields.append(pa.field(offset_id, pa.int64()))
+        schema = pa.schema(fields)
+        pt = pa.Table.from_pylist([], schema=schema)
+    else:
+        pt = pa.Table.from_pylist(rows)
+        pt = pt.rename_columns([*col_ids, offset_id])
+
+    return core.ArrayValue.from_pyarrow(pt, session=session)
 
 
 def _resolve_index_col(

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -663,7 +663,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // us + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         x_int = x.year() * 12 + x.month() - 1  # type: ignore
         first = y.year() * 12 + y.month() - 1  # type: ignore
         x_int_label = (
@@ -672,7 +672,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // n + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         x_int = x.year() * 4 + x.quarter() - 1  # type: ignore
         first = y.year() * 4 + y.quarter() - 1  # type: ignore
         x_int_label = (
@@ -681,7 +681,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // n + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         x_int = x.year()  # type: ignore
         first = y.year()  # type: ignore
         x_int_label = (
@@ -749,7 +749,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
             .cast(ibis_dtypes.Timestamp(timezone="UTC"))
             .cast(y.type())
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         one = ibis_types.literal(1)
         twelve = ibis_types.literal(12)
         first = y.year() * twelve + y.month() - one  # type: ignore
@@ -769,7 +769,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
             0,
         )
         x_label = next_month_date - ibis_api.interval(days=1)
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         one = ibis_types.literal(1)
         three = ibis_types.literal(3)
         four = ibis_types.literal(4)
@@ -792,7 +792,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
         )
 
         x_label = next_month_date - ibis_api.interval(days=1)
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         one = ibis_types.literal(1)
         first = y.year()  # type: ignore
         x = x * n + first  # type: ignore

bigframes/core/compile/sqlglot/aggregate_compiler.py

@@ -70,7 +70,5 @@ def compile_analytic(
             aggregate.arg.output_type,
         )
         return unary_compiler.compile(aggregate.op, column, window)
-    elif isinstance(aggregate, agg_expressions.BinaryAggregation):
-        raise NotImplementedError("binary analytic operations not yet supported")
     else:
         raise ValueError(f"Unexpected analytic operation: {aggregate}")

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

@@ -105,31 +105,6 @@ def _coerce_bool_to_int(typed_expr: TypedExpr) -> sge.Expression:
     return typed_expr.expr
 
 
-def _string_slice(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
-    # local name for each element in the array
-    el = sg.to_identifier("el")
-    # local name for the index in the array
-    slice_idx = sg.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)
-
-    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)
-
-
 def _array_slice(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
     # local name for each element in the array
     el = sg.to_identifier("el")

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

@@ -125,7 +125,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         x_int = sge.Paren(  # type: ignore
             this=sge.Add(
                 this=sge.Mul(
@@ -182,7 +182,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         x_int = sge.Paren(  # type: ignore
             this=sge.Add(
                 this=sge.Mul(
@@ -239,7 +239,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         x_int = sge.Extract(this=sge.Identifier(this="YEAR"), expression=x.expr)
         first = sge.Extract(this=sge.Identifier(this="YEAR"), expression=y.expr)
         return sge.Case(

bigframes/core/tree_properties.py

@@ -15,10 +15,13 @@
 
 import functools
 import itertools
-from typing import Callable, Dict, Optional, Sequence
+from typing import Callable, Dict, Optional, Sequence, TYPE_CHECKING
 
 import bigframes.core.nodes as nodes
 
+if TYPE_CHECKING:
+    import bigframes.session.execution_cache as execution_cache
+
 
 def is_trivially_executable(node: nodes.BigFrameNode) -> bool:
     if local_only(node):
@@ -65,7 +68,7 @@ def select_cache_target(
     root: nodes.BigFrameNode,
     min_complexity: float,
     max_complexity: float,
-    cache: dict[nodes.BigFrameNode, nodes.BigFrameNode],
+    cache: execution_cache.ExecutionCache,
     heuristic: Callable[[int, int], float],
 ) -> Optional[nodes.BigFrameNode]:
     """Take tree, and return candidate nodes with (# of occurences, post-caching planning complexity).
@@ -75,7 +78,7 @@ def select_cache_target(
 
     @functools.cache
     def _with_caching(subtree: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        return nodes.top_down(subtree, lambda x: cache.get(x, x))
+        return cache.subsitute_cached_subplans(subtree)
 
     def _combine_counts(
         left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int]
@@ -106,6 +109,7 @@ def _node_counts_inner(
     if len(node_counts) == 0:
         raise ValueError("node counts should be non-zero")
 
+    # for each considered node, calculate heuristic value, and return node with max value
     return max(
         node_counts.keys(),
         key=lambda node: heuristic(

bigframes/extensions/pandas/__init__.py

@@ -11,3 +11,15 @@
 # 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.
+
+"""
+BigQuery DataFrames automatically registers a pandas extenstion when imported.
+This allows you to use the power of the BigQuery engine with pandas objects
+directly.
+"""
+
+from bigframes.extensions.pandas.dataframe_accessor import (
+    PandasBigQueryDataFrameAccessor,
+)
+
+__all__ = ["PandasBigQueryDataFrameAccessor"]

bigframes/extensions/pandas/dataframe_accessor.py

@@ -18,10 +18,12 @@
 import pandas.api.extensions
 
 import bigframes.core.global_session as bf_session
+from bigframes.core.logging import log_adapter
 import bigframes.pandas as bpd
 
 
-class AIAccessor:
+@log_adapter.class_logger
+class PandasAIAccessor:
     """
     Pandas DataFrame accessor for BigQuery AI functions.
     """
@@ -101,7 +103,8 @@ def forecast(
 
 
 @pandas.api.extensions.register_dataframe_accessor("bigquery")
-class BigQueryDataFrameAccessor:
+@log_adapter.class_logger
+class PandasBigQueryDataFrameAccessor:
     """
     Pandas DataFrame accessor for BigQuery DataFrames functionality.
 
@@ -112,11 +115,11 @@ def __init__(self, pandas_obj: pandas.DataFrame):
         self._obj = pandas_obj
 
     @property
-    def ai(self) -> "AIAccessor":
+    def ai(self) -> "PandasAIAccessor":
         """
         Accessor for BigQuery AI functions.
         """
-        return AIAccessor(self._obj)
+        return PandasAIAccessor(self._obj)
 
     def sql_scalar(self, sql_template: str, *, output_dtype=None, session=None):
         """

bigframes/session/__init__.py

@@ -265,15 +265,20 @@ def __init__(
             metrics=self._metrics,
             publisher=self._publisher,
         )
+
+        labels = {}
+        if not self._strictly_ordered:
+            labels["bigframes-mode"] = "unordered"
+
         self._executor: executor.Executor = bq_caching_executor.BigQueryCachingExecutor(
             bqclient=self._clients_provider.bqclient,
             bqstoragereadclient=self._clients_provider.bqstoragereadclient,
             loader=self._loader,
             storage_manager=self._temp_storage_manager,
-            strictly_ordered=self._strictly_ordered,
             metrics=self._metrics,
             enable_polars_execution=context.enable_polars_execution,
             publisher=self._publisher,
+            labels=labels,
         )
 
     def __del__(self):