Merge branch 'main' into rf-vpc-fix

Author: jialuoo

bigframes/core/blocks.py

@@ -27,7 +27,6 @@
 import functools
 import itertools
 import random
-import textwrap
 import typing
 from typing import (
     Iterable,
@@ -54,16 +53,13 @@
 from bigframes.core import agg_expressions, local_data
 import bigframes.core as core
 import bigframes.core.agg_expressions as ex_types
-import bigframes.core.compile.googlesql as googlesql
 import bigframes.core.expression as ex
 import bigframes.core.expression as scalars
 import bigframes.core.guid as guid
 import bigframes.core.identifiers
 import bigframes.core.join_def as join_defs
 import bigframes.core.ordering as ordering
 import bigframes.core.pyarrow_utils as pyarrow_utils
-import bigframes.core.schema as bf_schema
-import bigframes.core.sql as sql
 import bigframes.core.utils as utils
 import bigframes.core.window_spec as windows
 import bigframes.dtypes
@@ -2776,14 +2772,6 @@ def _throw_if_null_index(self, opname: str):
             )
 
     def _get_rows_as_json_values(self) -> Block:
-        # We want to preserve any ordering currently present before turning to
-        # direct SQL manipulation. We will restore the ordering when we rebuild
-        # expression.
-        # TODO(shobs): Replace direct SQL manipulation by structured expression
-        # manipulation
-        expr, ordering_column_name = self.expr.promote_offsets()
-        expr_sql = self.session._executor.to_sql(expr)
-
         # Names of the columns to serialize for the row.
         # We will use the repr-eval pattern to serialize a value here and
         # deserialize in the cloud function. Let's make sure that would work.
@@ -2799,93 +2787,44 @@ def _get_rows_as_json_values(self) -> Block:
                 )
 
             column_names.append(serialized_column_name)
-        column_names_csv = sql.csv(map(sql.simple_literal, column_names))
-
-        # index columns count
-        index_columns_count = len(self.index_columns)
 
         # column references to form the array of values for the row
         column_types = list(self.index.dtypes) + list(self.dtypes)
         column_references = []
         for type_, col in zip(column_types, self.expr.column_ids):
-            if isinstance(type_, pd.ArrowDtype) and pa.types.is_binary(
-                type_.pyarrow_dtype
-            ):
-                column_references.append(sql.to_json_string(col))
+            if type_ == bigframes.dtypes.BYTES_DTYPE:
+                column_references.append(ops.ToJSONString().as_expr(col))
+            elif type_ == bigframes.dtypes.BOOL_DTYPE:
+                # cast operator produces True/False, but function template expects lower case
+                column_references.append(
+                    ops.lower_op.as_expr(
+                        ops.AsTypeOp(bigframes.dtypes.STRING_DTYPE).as_expr(col)
+                    )
+                )
             else:
-                column_references.append(sql.cast_as_string(col))
-
-        column_references_csv = sql.csv(column_references)
-
-        # types of the columns to serialize for the row
-        column_types_csv = sql.csv(
-            [sql.simple_literal(str(typ)) for typ in column_types]
-        )
+                column_references.append(
+                    ops.AsTypeOp(bigframes.dtypes.STRING_DTYPE).as_expr(col)
+                )
 
         # row dtype to use for deserializing the row as pandas series
         pandas_row_dtype = bigframes.dtypes.lcd_type(*column_types)
         if pandas_row_dtype is None:
             pandas_row_dtype = "object"
-        pandas_row_dtype = sql.simple_literal(str(pandas_row_dtype))
-
-        # create a json column representing row through SQL manipulation
-        row_json_column_name = guid.generate_guid()
-        select_columns = (
-            [ordering_column_name] + list(self.index_columns) + [row_json_column_name]
-        )
-        select_columns_csv = sql.csv(
-            [googlesql.identifier(col) for col in select_columns]
-        )
-        json_sql = f"""\
-With T0 AS (
-{textwrap.indent(expr_sql, "    ")}
-),
-T1 AS (
-    SELECT *,
-           TO_JSON_STRING(JSON_OBJECT(
-               "names", [{column_names_csv}],
-               "types", [{column_types_csv}],
-               "values", [{column_references_csv}],
-               "indexlength", {index_columns_count},
-               "dtype", {pandas_row_dtype}
-           )) AS {googlesql.identifier(row_json_column_name)} FROM T0
-)
-SELECT {select_columns_csv} FROM T1
-"""
-        # The only ways this code is used is through df.apply(axis=1) cope path
-        destination, query_job = self.session._loader._query_to_destination(
-            json_sql, cluster_candidates=[ordering_column_name]
-        )
-        if not destination:
-            raise ValueError(f"Query job {query_job} did not produce result table")
-
-        new_schema = (
-            self.expr.schema.select([*self.index_columns])
-            .append(
-                bf_schema.SchemaItem(
-                    row_json_column_name, bigframes.dtypes.STRING_DTYPE
-                )
-            )
-            .append(
-                bf_schema.SchemaItem(ordering_column_name, bigframes.dtypes.INT_DTYPE)
-            )
-        )
+        pandas_row_dtype = str(pandas_row_dtype)
 
-        dest_table = self.session.bqclient.get_table(destination)
-        expr = core.ArrayValue.from_table(
-            dest_table,
-            schema=new_schema,
-            session=self.session,
-            offsets_col=ordering_column_name,
-            n_rows=dest_table.num_rows,
-        ).drop_columns([ordering_column_name])
-        block = Block(
-            expr,
-            index_columns=self.index_columns,
-            column_labels=[row_json_column_name],
-            index_labels=self._index_labels,
+        struct_op = ops.StructOp(
+            column_names=("names", "types", "values", "indexlength", "dtype")
         )
-        return block
+        names_val = ex.const(tuple(column_names))
+        types_val = ex.const(tuple(map(str, column_types)))
+        values_val = ops.ToArrayOp().as_expr(*column_references)
+        indexlength_val = ex.const(len(self.index_columns))
+        dtype_val = ex.const(str(pandas_row_dtype))
+        struct_expr = struct_op.as_expr(
+            names_val, types_val, values_val, indexlength_val, dtype_val
+        )
+        block, col_id = self.project_expr(ops.ToJSONString().as_expr(struct_expr))
+        return block.select_column(col_id)
 
 
 class BlockIndexProperties:

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1301,8 +1301,8 @@ def parse_json_op_impl(x: ibis_types.Value, op: ops.ParseJSON):
 
 
 @scalar_op_compiler.register_unary_op(ops.ToJSONString)
-def to_json_string_op_impl(json_obj: ibis_types.Value):
-    return to_json_string(json_obj=json_obj)
+def to_json_string_op_impl(x: ibis_types.Value):
+    return to_json_string(value=x)
 
 
 @scalar_op_compiler.register_unary_op(ops.JSONValue, pass_op=True)
@@ -2069,9 +2069,9 @@ def json_extract_string_array(  # type: ignore[empty-body]
 
 @ibis_udf.scalar.builtin(name="to_json_string")
 def to_json_string(  # type: ignore[empty-body]
-    json_obj: ibis_dtypes.JSON,
+    value,
 ) -> ibis_dtypes.String:
-    """Convert JSON to STRING."""
+    """Convert value to JSON-formatted string."""
 
 
 @ibis_udf.scalar.builtin(name="json_value")

bigframes/core/compile/ibis_types.py

@@ -386,10 +386,6 @@ def literal_to_ibis_scalar(
     ibis_dtype = bigframes_dtype_to_ibis_dtype(force_dtype) if force_dtype else None
 
     if pd.api.types.is_list_like(literal):
-        if validate:
-            raise ValueError(
-                f"List types can't be stored in BigQuery DataFrames. {constants.FEEDBACK_LINK}"
-            )
         # "correct" way would be to use ibis.array, but this produces invalid BQ SQL syntax
         return tuple(literal)
 

bigframes/core/groupby/dataframe_group_by.py

@@ -461,23 +461,19 @@ def expanding(self, min_periods: int = 1) -> windows.Window:
 
     def agg(self, func=None, **kwargs) -> typing.Union[df.DataFrame, series.Series]:
         if func:
-            if isinstance(func, str):
-                return self.size() if func == "size" else self._agg_string(func)
-            elif utils.is_dict_like(func):
+            if utils.is_dict_like(func):
                 return self._agg_dict(func)
             elif utils.is_list_like(func):
                 return self._agg_list(func)
             else:
-                raise NotImplementedError(
-                    f"Aggregate with {func} not supported. {constants.FEEDBACK_LINK}"
-                )
+                return self.size() if func == "size" else self._agg_func(func)
         else:
             return self._agg_named(**kwargs)
 
-    def _agg_string(self, func: str) -> df.DataFrame:
+    def _agg_func(self, func) -> df.DataFrame:
         ids, labels = self._aggregated_columns()
         aggregations = [
-            aggs.agg(col_id, agg_ops.lookup_agg_func(func)) for col_id in ids
+            aggs.agg(col_id, agg_ops.lookup_agg_func(func)[0]) for col_id in ids
         ]
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -500,7 +496,7 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
                 funcs_for_id if utils.is_list_like(funcs_for_id) else [funcs_for_id]
             )
             for f in func_list:
-                aggregations.append(aggs.agg(col_id, agg_ops.lookup_agg_func(f)))
+                aggregations.append(aggs.agg(col_id, agg_ops.lookup_agg_func(f)[0]))
                 column_labels.append(label)
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -525,19 +521,23 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
     def _agg_list(self, func: typing.Sequence) -> df.DataFrame:
         ids, labels = self._aggregated_columns()
         aggregations = [
-            aggs.agg(col_id, agg_ops.lookup_agg_func(f)) for col_id in ids for f in func
+            aggs.agg(col_id, agg_ops.lookup_agg_func(f)[0])
+            for col_id in ids
+            for f in func
         ]
 
         if self._block.column_labels.nlevels > 1:
             # Restructure MultiIndex for proper format: (idx1, idx2, func)
             # rather than ((idx1, idx2), func).
             column_labels = [
-                tuple(label) + (f,)
+                tuple(label) + (agg_ops.lookup_agg_func(f)[1],)
                 for label in labels.to_frame(index=False).to_numpy()
                 for f in func
             ]
         else:  # Single-level index
-            column_labels = [(label, f) for label in labels for f in func]
+            column_labels = [
+                (label, agg_ops.lookup_agg_func(f)[1]) for label in labels for f in func
+            ]
 
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -563,7 +563,7 @@ def _agg_named(self, **kwargs) -> df.DataFrame:
             if not isinstance(v, tuple) or (len(v) != 2):
                 raise TypeError("kwargs values must be 2-tuples of column, aggfunc")
             col_id = self._resolve_label(v[0])
-            aggregations.append(aggs.agg(col_id, agg_ops.lookup_agg_func(v[1])))
+            aggregations.append(aggs.agg(col_id, agg_ops.lookup_agg_func(v[1])[0]))
             column_labels.append(k)
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,

bigframes/core/groupby/series_group_by.py

@@ -216,18 +216,17 @@ def prod(self, *args) -> series.Series:
 
     def agg(self, func=None) -> typing.Union[df.DataFrame, series.Series]:
         column_names: list[str] = []
-        if isinstance(func, str):
-            aggregations = [aggs.agg(self._value_column, agg_ops.lookup_agg_func(func))]
-            column_names = [func]
-        elif utils.is_list_like(func):
-            aggregations = [
-                aggs.agg(self._value_column, agg_ops.lookup_agg_func(f)) for f in func
-            ]
-            column_names = list(func)
-        else:
+        if utils.is_dict_like(func):
             raise NotImplementedError(
                 f"Aggregate with {func} not supported. {constants.FEEDBACK_LINK}"
             )
+        if not utils.is_list_like(func):
+            func = [func]
+
+        aggregations = [
+            aggs.agg(self._value_column, agg_ops.lookup_agg_func(f)[0]) for f in func
+        ]
+        column_names = [agg_ops.lookup_agg_func(f)[1] for f in func]
 
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,

bigframes/core/log_adapter.py

@@ -149,49 +149,61 @@ def wrap(cls):
     return wrap(decorated_cls)
 
 
-def method_logger(method, /, *, custom_base_name: Optional[str] = None):
+def method_logger(method=None, /, *, custom_base_name: Optional[str] = None):
     """Decorator that adds logging functionality to a method."""
 
-    @functools.wraps(method)
-    def wrapper(*args, **kwargs):
-        api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
-        if custom_base_name is None:
-            qualname_parts = getattr(method, "__qualname__", method.__name__).split(".")
-            class_name = qualname_parts[-2] if len(qualname_parts) > 1 else ""
-            base_name = (
-                class_name if class_name else "_".join(method.__module__.split(".")[1:])
-            )
-        else:
-            base_name = custom_base_name
-
-        full_method_name = f"{base_name.lower()}-{api_method_name}"
-        # Track directly called methods
-        if len(_call_stack) == 0:
-            add_api_method(full_method_name)
-
-        _call_stack.append(full_method_name)
-
-        try:
-            return method(*args, **kwargs)
-        except (NotImplementedError, TypeError) as e:
-            # Log method parameters that are implemented in pandas but either missing (TypeError)
-            # or not fully supported (NotImplementedError) in BigFrames.
-            # Logging is currently supported only when we can access the bqclient through
-            # _block.session.bqclient.
-            if len(_call_stack) == 1:
-                submit_pandas_labels(
-                    _get_bq_client(*args, **kwargs),
-                    base_name,
-                    api_method_name,
-                    args,
-                    kwargs,
-                    task=PANDAS_PARAM_TRACKING_TASK,
+    def outer_wrapper(method):
+        @functools.wraps(method)
+        def wrapper(*args, **kwargs):
+            api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
+            if custom_base_name is None:
+                qualname_parts = getattr(method, "__qualname__", method.__name__).split(
+                    "."
+                )
+                class_name = qualname_parts[-2] if len(qualname_parts) > 1 else ""
+                base_name = (
+                    class_name
+                    if class_name
+                    else "_".join(method.__module__.split(".")[1:])
                 )
-            raise e
-        finally:
-            _call_stack.pop()
+            else:
+                base_name = custom_base_name
 
-    return wrapper
+            full_method_name = f"{base_name.lower()}-{api_method_name}"
+            # Track directly called methods
+            if len(_call_stack) == 0:
+                add_api_method(full_method_name)
+
+            _call_stack.append(full_method_name)
+
+            try:
+                return method(*args, **kwargs)
+            except (NotImplementedError, TypeError) as e:
+                # Log method parameters that are implemented in pandas but either missing (TypeError)
+                # or not fully supported (NotImplementedError) in BigFrames.
+                # Logging is currently supported only when we can access the bqclient through
+                # _block.session.bqclient.
+                if len(_call_stack) == 1:
+                    submit_pandas_labels(
+                        _get_bq_client(*args, **kwargs),
+                        base_name,
+                        api_method_name,
+                        args,
+                        kwargs,
+                        task=PANDAS_PARAM_TRACKING_TASK,
+                    )
+                raise e
+            finally:
+                _call_stack.pop()
+
+        return wrapper
+
+    if method is None:
+        # Called with parentheses
+        return outer_wrapper
+
+    # Called without parentheses
+    return outer_wrapper(method)
 
 
 def property_logger(prop):