Merge remote-tracking branch 'github/main' into align_df_series

Author: TrevorBergeron

.kokoro/continuous/e2e.cfg

@@ -3,7 +3,7 @@
 # Only run this nox session.
 env_vars: {
     key: "NOX_SESSION"
-    value: "unit_prerelease system_prerelease system_noextras e2e notebook"
+    value: "e2e doctest notebook unit_prerelease system_prerelease system_noextras"
 }
 
 env_vars: {

.kokoro/presubmit/e2e.cfg

@@ -3,7 +3,7 @@
 # Only run this nox session.
 env_vars: {
     key: "NOX_SESSION"
-    value: "unit_prerelease system_prerelease system_noextras e2e notebook"
+    value: "e2e doctest notebook unit_prerelease system_prerelease system_noextras"
 }
 
 env_vars: {

bigframes/constants.py

@@ -12,19 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
-
 """Constants used across BigQuery DataFrames.
 
 This module should not depend on any others in the package.
 """
 
-FEEDBACK_LINK = (
-    "Share your usecase with the BigQuery DataFrames team at the "
-    "https://bit.ly/bigframes-feedback survey."
-)
+import datetime
+
+import bigframes_vendored.constants
 
-ABSTRACT_METHOD_ERROR_MESSAGE = f"Abstract method. You have likely encountered a bug. Please share this stacktrace and how you reached it with the BigQuery DataFrames team. {FEEDBACK_LINK}"
+FEEDBACK_LINK = bigframes_vendored.constants.FEEDBACK_LINK
+ABSTRACT_METHOD_ERROR_MESSAGE = (
+    bigframes_vendored.constants.ABSTRACT_METHOD_ERROR_MESSAGE
+)
 
 DEFAULT_EXPIRATION = datetime.timedelta(days=7)
 

bigframes/core/blocks.py

@@ -1036,6 +1036,34 @@ def aggregate_all_and_stack(
                 index_labels=self.index.names,
             )
 
+    def aggregate_size(
+        self,
+        by_column_ids: typing.Sequence[str] = (),
+        *,
+        dropna: bool = True,
+    ):
+        """Returns a block object to compute the size(s) of groups."""
+        agg_specs = [
+            (ex.NullaryAggregation(agg_ops.SizeOp()), guid.generate_guid()),
+        ]
+        output_col_ids = [agg_spec[1] for agg_spec in agg_specs]
+        result_expr = self.expr.aggregate(agg_specs, by_column_ids, dropna=dropna)
+        names: typing.List[Label] = []
+        for by_col_id in by_column_ids:
+            if by_col_id in self.value_columns:
+                names.append(self.col_id_to_label[by_col_id])
+            else:
+                names.append(self.col_id_to_index_name[by_col_id])
+        return (
+            Block(
+                result_expr,
+                index_columns=by_column_ids,
+                column_labels=["size"],
+                index_labels=names,
+            ),
+            output_col_ids,
+        )
+
     def select_column(self, id: str) -> Block:
         return self.select_columns([id])
 

bigframes/core/compile/aggregate_compiler.py

@@ -36,6 +36,8 @@ def compile_aggregate(
     bindings: typing.Dict[str, ibis_types.Value],
     order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
+    if isinstance(aggregate, ex.NullaryAggregation):
+        return compile_nullary_agg(aggregate.op)
     if isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         if aggregate.op.can_order_by:
@@ -55,7 +57,9 @@ def compile_analytic(
     window: window_spec.WindowSpec,
     bindings: typing.Dict[str, ibis_types.Value],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, ex.UnaryAggregation):
+    if isinstance(aggregate, ex.NullaryAggregation):
+        return compile_nullary_agg(aggregate.op, window)
+    elif isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         return compile_unary_agg(aggregate.op, input, window)
     elif isinstance(aggregate, ex.BinaryAggregation):
@@ -93,6 +97,14 @@ def compile_ordered_unary_agg(
     raise ValueError(f"Can't compile unrecognized operation: {op}")
 
 
+@functools.singledispatch
+def compile_nullary_agg(
+    op: agg_ops.WindowOp,
+    window: Optional[window_spec.WindowSpec] = None,
+) -> ibis_types.Value:
+    raise ValueError(f"Can't compile unrecognized operation: {op}")
+
+
 def numeric_op(operation):
     @functools.wraps(operation)
     def constrained_op(
@@ -118,6 +130,11 @@ def constrained_op(
 ### Specific Op implementations Below
 
 
+@compile_nullary_agg.register
+def _(op: agg_ops.SizeOp, window=None) -> ibis_types.NumericValue:
+    return _apply_window_if_present(vendored_ibis_ops.count(1), window)
+
+
 @compile_unary_agg.register
 @numeric_op
 def _(

bigframes/core/compile/scalar_op_compiler.py

@@ -856,11 +856,12 @@ def to_timestamp_op_impl(x: ibis_types.Value, op: ops.ToTimestampOp):
 
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)
 def remote_function_op_impl(x: ibis_types.Value, op: ops.RemoteFunctionOp):
-    if not hasattr(op.func, "bigframes_remote_function"):
+    ibis_node = getattr(op.func, "ibis_node", None)
+    if ibis_node is None:
         raise TypeError(
             f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}"
         )
-    x_transformed = op.func(x)
+    x_transformed = ibis_node(x)
     if not op.apply_on_null:
         x_transformed = ibis.case().when(x.isnull(), x).else_(x_transformed).end()
     return x_transformed
@@ -1342,11 +1343,12 @@ def minimum_impl(
 def binary_remote_function_op_impl(
     x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp
 ):
-    if not hasattr(op.func, "bigframes_remote_function"):
+    ibis_node = getattr(op.func, "ibis_node", None)
+    if ibis_node is None:
         raise TypeError(
             f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}"
         )
-    x_transformed = op.func(x, y)
+    x_transformed = ibis_node(x, y)
     return x_transformed
 
 

bigframes/core/expression.py

@@ -46,6 +46,16 @@ def output_type(
         ...
 
 
+@dataclasses.dataclass(frozen=True)
+class NullaryAggregation(Aggregation):
+    op: agg_ops.NullaryWindowOp = dataclasses.field()
+
+    def output_type(
+        self, input_types: dict[str, bigframes.dtypes.Dtype]
+    ) -> dtypes.ExpressionType:
+        return self.op.output_type()
+
+
 @dataclasses.dataclass(frozen=True)
 class UnaryAggregation(Aggregation):
     op: agg_ops.UnaryWindowOp = dataclasses.field()

bigframes/core/groupby/__init__.py

@@ -104,6 +104,20 @@ def __getitem__(
                 dropna=self._dropna,
             )
 
+    def size(self) -> typing.Union[df.DataFrame, series.Series]:
+        agg_block, _ = self._block.aggregate_size(
+            by_column_ids=self._by_col_ids,
+            dropna=self._dropna,
+        )
+        agg_block = agg_block.with_column_labels(pd.Index(["size"]))
+        dataframe = df.DataFrame(agg_block)
+
+        if self._as_index:
+            series = dataframe["size"]
+            return series.rename(None)
+        else:
+            return self._convert_index(dataframe)
+
     def sum(self, numeric_only: bool = False, *args) -> df.DataFrame:
         if not numeric_only:
             self._raise_on_non_numeric("sum")
@@ -520,6 +534,13 @@ def std(self, *args, **kwargs) -> series.Series:
     def var(self, *args, **kwargs) -> series.Series:
         return self._aggregate(agg_ops.var_op)
 
+    def size(self) -> series.Series:
+        agg_block, _ = self._block.aggregate_size(
+            by_column_ids=self._by_col_ids,
+            dropna=self._dropna,
+        )
+        return series.Series(agg_block, name=self._value_name)
+
     def skew(self, *args, **kwargs) -> series.Series:
         block = block_ops.skew(self._block, [self._value_column], self._by_col_ids)
         return series.Series(block)

bigframes/core/log_adapter.py

@@ -99,9 +99,12 @@ def add_api_method(api_method_name):
         _api_methods = _api_methods[:MAX_LABELS_COUNT]
 
 
-def get_and_reset_api_methods():
+def get_and_reset_api_methods(dry_run: bool = False):
     global _lock
     with _lock:
         previous_api_methods = list(_api_methods)
-        _api_methods.clear()
+
+        # dry_run might not make a job resource, so only reset the log on real queries.
+        if not dry_run:
+            _api_methods.clear()
     return previous_api_methods

bigframes/dataframe.py

@@ -400,11 +400,12 @@ def memory_usage(self, index: bool = True):
         column_sizes = self.dtypes.map(
             lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows
         )
-        if index:
+        if index and self._has_index:
             index_size = pandas.Series([self.index._memory_usage()], index=["Index"])
             column_sizes = pandas.concat([index_size, column_sizes])
         return column_sizes
 
+    @requires_index
     def info(
         self,
         verbose: Optional[bool] = None,
@@ -778,7 +779,7 @@ def _apply_series_binop_axis_0(
             block = block.drop_columns([get_column_left[column_id]])
 
         block = block.drop_columns([series_col])
-        block = block.with_index_labels(self.index.names)
+        block = block.with_index_labels(self._block.index.names)
         return DataFrame(block)
 
     def _apply_series_binop_axis_1(
@@ -914,6 +915,11 @@ def __ne__(self, other) -> DataFrame:  # type: ignore
 
     __ne__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__ne__)
 
+    def __invert__(self) -> DataFrame:
+        return self._apply_unary_op(ops.invert_op)
+
+    __invert__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__invert__)
+
     def le(self, other: typing.Any, axis: str | int = "columns") -> DataFrame:
         return self._apply_binop(other, ops.le_op, axis=axis)
 
@@ -1622,7 +1628,7 @@ def _assign_series_join_on_index(
             # Update case, remove after copying into columns
             block = block.drop_columns([source_column])
 
-        return DataFrame(block.with_index_labels(self.index.names))
+        return DataFrame(block.with_index_labels(self._block.index.names))
 
     def reset_index(self, *, drop: bool = False) -> DataFrame:
         block = self._block.reset_index(drop)
@@ -2917,7 +2923,9 @@ def to_csv(
             field_delimiter=sep,
             header=header,
         )
-        _, query_job = self._block.expr.session._start_query(export_data_statement)
+        _, query_job = self._block.expr.session._start_query(
+            export_data_statement, api_name="dataframe-to_csv"
+        )
         self._set_internal_query_job(query_job)
 
     def to_json(
@@ -2959,7 +2967,9 @@ def to_json(
             format="JSON",
             export_options={},
         )
-        _, query_job = self._block.expr.session._start_query(export_data_statement)
+        _, query_job = self._block.expr.session._start_query(
+            export_data_statement, api_name="dataframe-to_json"
+        )
         self._set_internal_query_job(query_job)
 
     def to_gbq(
@@ -3091,7 +3101,9 @@ def to_parquet(
             format="PARQUET",
             export_options=export_options,
         )
-        _, query_job = self._block.expr.session._start_query(export_data_statement)
+        _, query_job = self._block.expr.session._start_query(
+            export_data_statement, api_name="dataframe-to_parquet"
+        )
         self._set_internal_query_job(query_job)
 
     def to_dict(
@@ -3294,7 +3306,7 @@ def _prepare_export(
         array_value = self._block.expr
 
         new_col_labels, new_idx_labels = utils.get_standardized_ids(
-            self._block.column_labels, self.index.names
+            self._block.column_labels, self._block.index.names
         )
 
         columns = list(self._block.value_columns)
@@ -3331,7 +3343,7 @@ def _run_io_query(
         session = self._block.expr.session
         self._optimize_query_complexity()
         export_array, id_overrides = self._prepare_export(
-            index=index, ordering_id=ordering_id
+            index=index and self._has_index, ordering_id=ordering_id
         )
 
         _, query_job = session._execute(

bigframes/exceptions.py

@@ -41,3 +41,7 @@ class PreviewWarning(Warning):
 
 class NullIndexError(ValueError):
     """Object has no index."""
+
+
+class TimeTravelDisabledWarning(Warning):
+    """A query was reattempted without time travel."""