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

Author: TrevorBergeron

CHANGELOG.md

@@ -4,6 +4,25 @@
 
 [1]: https://pypi.org/project/bigframes/#history
 
+## [2.10.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.9.0...v2.10.0) (2025-07-08)
+
+
+### Features
+
+* `df.to_pandas_batches()` returns one empty DataFrame if `df` is empty ([#1878](https://github.com/googleapis/python-bigquery-dataframes/issues/1878)) ([e43d15d](https://github.com/googleapis/python-bigquery-dataframes/commit/e43d15d535d6d5fd73c33967271f3591c41dffb3))
+* Add filter pushdown to hybrid engine ([#1871](https://github.com/googleapis/python-bigquery-dataframes/issues/1871)) ([6454aff](https://github.com/googleapis/python-bigquery-dataframes/commit/6454aff726dee791acbac98f893075ee5ee6d9a1))
+* Add simple stats support to hybrid local pushdown ([#1873](https://github.com/googleapis/python-bigquery-dataframes/issues/1873)) ([8715105](https://github.com/googleapis/python-bigquery-dataframes/commit/8715105239216bffe899ddcbb15805f2e3063af4))
+
+
+### Bug Fixes
+
+* Fix issues where duration type returned as int ([#1875](https://github.com/googleapis/python-bigquery-dataframes/issues/1875)) ([f30f750](https://github.com/googleapis/python-bigquery-dataframes/commit/f30f75053a6966abd1a6a644c23efb86b2ac568d))
+
+
+### Documentation
+
+* Update gsutil commands to gcloud commands ([#1876](https://github.com/googleapis/python-bigquery-dataframes/issues/1876)) ([c289f70](https://github.com/googleapis/python-bigquery-dataframes/commit/c289f7061320ec6d9de099cab2416cc9f289baac))
+
 ## [2.9.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.8.0...v2.9.0) (2025-06-30)
 
 

bigframes/core/blocks.py

@@ -29,7 +29,17 @@
 import random
 import textwrap
 import typing
-from typing import Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import (
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 import warnings
 
 import bigframes_vendored.constants as constants
@@ -87,14 +97,22 @@
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]
 
 
-class BlockHolder(typing.Protocol):
+@dataclasses.dataclass
+class PandasBatches(Iterator[pd.DataFrame]):
     """Interface for mutable objects with state represented by a block value object."""
 
-    def _set_block(self, block: Block):
-        """Set the underlying block value of the object"""
+    def __init__(
+        self, pandas_batches: Iterator[pd.DataFrame], total_rows: Optional[int] = 0
+    ):
+        self._dataframes: Iterator[pd.DataFrame] = pandas_batches
+        self._total_rows: Optional[int] = total_rows
+
+    @property
+    def total_rows(self) -> Optional[int]:
+        return self._total_rows
 
-    def _get_block(self) -> Block:
-        """Get the underlying block value of the object"""
+    def __next__(self) -> pd.DataFrame:
+        return next(self._dataframes)
 
 
 @dataclasses.dataclass()
@@ -599,8 +617,7 @@ def try_peek(
                 self.expr, n, use_explicit_destination=allow_large_results
             )
             df = result.to_pandas()
-            self._copy_index_to_pandas(df)
-            return df
+            return self._copy_index_to_pandas(df)
         else:
             return None
 
@@ -609,8 +626,7 @@ def to_pandas_batches(
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         allow_large_results: Optional[bool] = None,
-        squeeze: Optional[bool] = False,
-    ):
+    ) -> Iterator[pd.DataFrame]:
         """Download results one message at a time.
 
         page_size and max_results determine the size and number of batches,
@@ -621,43 +637,43 @@ def to_pandas_batches(
             use_explicit_destination=allow_large_results,
         )
 
-        total_batches = 0
-        for df in execute_result.to_pandas_batches(
-            page_size=page_size, max_results=max_results
-        ):
-            total_batches += 1
-            self._copy_index_to_pandas(df)
-            if squeeze:
-                yield df.squeeze(axis=1)
-            else:
-                yield df
-
         # To reduce the number of edge cases to consider when working with the
         # results of this, always return at least one DataFrame. See:
         # b/428918844.
-        if total_batches == 0:
-            df = pd.DataFrame(
-                {
-                    col: pd.Series([], dtype=self.expr.get_column_type(col))
-                    for col in itertools.chain(self.value_columns, self.index_columns)
-                }
-            )
-            self._copy_index_to_pandas(df)
-            yield df
+        empty_val = pd.DataFrame(
+            {
+                col: pd.Series([], dtype=self.expr.get_column_type(col))
+                for col in itertools.chain(self.value_columns, self.index_columns)
+            }
+        )
+        dfs = map(
+            lambda a: a[0],
+            itertools.zip_longest(
+                execute_result.to_pandas_batches(page_size, max_results),
+                [0],
+                fillvalue=empty_val,
+            ),
+        )
+        dfs = iter(map(self._copy_index_to_pandas, dfs))
 
-    def _copy_index_to_pandas(self, df: pd.DataFrame):
-        """Set the index on pandas DataFrame to match this block.
+        total_rows = execute_result.total_rows
+        if (total_rows is not None) and (max_results is not None):
+            total_rows = min(total_rows, max_results)
 
-        Warning: This method modifies ``df`` inplace.
-        """
+        return PandasBatches(dfs, total_rows)
+
+    def _copy_index_to_pandas(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Set the index on pandas DataFrame to match this block."""
         # Note: If BigQuery DataFrame has null index, a default one will be created for the local materialization.
+        new_df = df.copy()
         if len(self.index_columns) > 0:
-            df.set_index(list(self.index_columns), inplace=True)
+            new_df.set_index(list(self.index_columns), inplace=True)
             # Pandas names is annotated as list[str] rather than the more
             # general Sequence[Label] that BigQuery DataFrames has.
             # See: https://github.com/pandas-dev/pandas-stubs/issues/804
-            df.index.names = self.index.names  # type: ignore
-        df.columns = self.column_labels
+            new_df.index.names = self.index.names  # type: ignore
+        new_df.columns = self.column_labels
+        return new_df
 
     def _materialize_local(
         self, materialize_options: MaterializationOptions = MaterializationOptions()
@@ -724,9 +740,7 @@ def _materialize_local(
             )
         else:
             df = execute_result.to_pandas()
-            self._copy_index_to_pandas(df)
-
-        return df, execute_result.query_job
+            return self._copy_index_to_pandas(df), execute_result.query_job
 
     def _downsample(
         self, total_rows: int, sampling_method: str, fraction: float, random_state
@@ -1591,8 +1605,7 @@ def retrieve_repr_request_results(
         row_count = self.session._executor.execute(self.expr.row_count()).to_py_scalar()
 
         head_df = head_result.to_pandas()
-        self._copy_index_to_pandas(head_df)
-        return head_df, row_count, head_result.query_job
+        return self._copy_index_to_pandas(head_df), row_count, head_result.query_job
 
     def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]:
         expr, result_id = self._expr.promote_offsets()

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

@@ -42,3 +42,8 @@ def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
 @BINARY_OP_REGISTRATION.register(ops.ge_op)
 def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.GTE(this=left.expr, expression=right.expr)
+
+
+@BINARY_OP_REGISTRATION.register(ops.JSONSet)
+def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_SET", left.expr, sge.convert(op.json_path), right.expr)

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

@@ -70,3 +70,49 @@ def _(op: ops.ArraySliceOp, expr: TypedExpr) -> sge.Expression:
     )
 
     return sge.array(selected_elements)
+
+
+# JSON Ops
+@UNARY_OP_REGISTRATION.register(ops.JSONExtract)
+def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_EXTRACT", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONExtractArray)
+def _(op: ops.JSONExtractArray, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_EXTRACT_ARRAY", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONExtractStringArray)
+def _(op: ops.JSONExtractStringArray, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_EXTRACT_STRING_ARRAY", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONQuery)
+def _(op: ops.JSONQuery, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_QUERY", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONQueryArray)
+def _(op: ops.JSONQueryArray, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_QUERY_ARRAY", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONValue)
+def _(op: ops.JSONValue, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_VALUE", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.JSONValueArray)
+def _(op: ops.JSONValueArray, expr: TypedExpr) -> sge.Expression:
+    return sge.func("JSON_VALUE_ARRAY", expr.expr, sge.convert(op.json_path))
+
+
+@UNARY_OP_REGISTRATION.register(ops.ParseJSON)
+def _(op: ops.ParseJSON, expr: TypedExpr) -> sge.Expression:
+    return sge.func("PARSE_JSON", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.ToJSONString)
+def _(op: ops.ToJSONString, expr: TypedExpr) -> sge.Expression:
+    return sge.func("TO_JSON_STRING", expr.expr)

bigframes/core/nodes.py

@@ -161,7 +161,7 @@ def is_noop(self) -> bool:
         return (
             ((not self.start) or (self.start == 0))
             and (self.step == 1)
-            and ((self.stop is None) or (self.stop == self.row_count))
+            and ((self.stop is None) or (self.stop == self.child.row_count))
         )
 
     @property

bigframes/series.py

@@ -648,13 +648,12 @@ def to_pandas_batches(
                 form the original Series. Results stream from bigquery,
                 see https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.table.RowIterator#google_cloud_bigquery_table_RowIterator_to_arrow_iterable
         """
-        df = self._block.to_pandas_batches(
+        batches = self._block.to_pandas_batches(
             page_size=page_size,
             max_results=max_results,
             allow_large_results=allow_large_results,
-            squeeze=True,
         )
-        return df
+        return map(lambda df: cast(pandas.Series, df.squeeze(1)), batches)
 
     def _compute_dry_run(self) -> bigquery.QueryJob:
         _, query_job = self._block._compute_dry_run((self._value_column,))

bigframes/session/bq_caching_executor.py

@@ -100,6 +100,7 @@ def cache_results_table(
         original_root: nodes.BigFrameNode,
         table: bigquery.Table,
         ordering: order.RowOrdering,
+        num_rows: Optional[int] = None,
     ):
         # Assumption: GBQ cached table uses field name as bq column name
         scan_list = nodes.ScanList(
@@ -112,7 +113,7 @@ def cache_results_table(
             source=nodes.BigqueryDataSource(
                 nodes.GbqTable.from_table(table),
                 ordering=ordering,
-                n_rows=table.num_rows,
+                n_rows=num_rows,
             ),
             scan_list=scan_list,
             table_session=original_root.session,
@@ -468,14 +469,16 @@ def _cache_with_cluster_cols(
                 plan, sort_rows=False, materialize_all_order_keys=True
             )
         )
-        tmp_table_ref = self._sql_as_cached_temp_table(
+        tmp_table_ref, num_rows = self._sql_as_cached_temp_table(
             compiled.sql,
             compiled.sql_schema,
             cluster_cols=bq_io.select_cluster_cols(compiled.sql_schema, cluster_cols),
         )
         tmp_table = self.bqclient.get_table(tmp_table_ref)
         assert compiled.row_order is not None
-        self.cache.cache_results_table(array_value.node, tmp_table, compiled.row_order)
+        self.cache.cache_results_table(
+            array_value.node, tmp_table, compiled.row_order, num_rows=num_rows
+        )
 
     def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
         """Executes the query and uses the resulting table to rewrite future executions."""
@@ -487,14 +490,16 @@ def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
                 sort_rows=False,
             )
         )
-        tmp_table_ref = self._sql_as_cached_temp_table(
+        tmp_table_ref, num_rows = self._sql_as_cached_temp_table(
             compiled.sql,
             compiled.sql_schema,
             cluster_cols=[offset_column],
         )
         tmp_table = self.bqclient.get_table(tmp_table_ref)
         assert compiled.row_order is not None
-        self.cache.cache_results_table(array_value.node, tmp_table, compiled.row_order)
+        self.cache.cache_results_table(
+            array_value.node, tmp_table, compiled.row_order, num_rows=num_rows
+        )
 
     def _cache_with_session_awareness(
         self,
@@ -552,7 +557,7 @@ def _sql_as_cached_temp_table(
         sql: str,
         schema: Sequence[bigquery.SchemaField],
         cluster_cols: Sequence[str],
-    ) -> bigquery.TableReference:
+    ) -> tuple[bigquery.TableReference, Optional[int]]:
         assert len(cluster_cols) <= _MAX_CLUSTER_COLUMNS
         temp_table = self.storage_manager.create_temp_table(schema, cluster_cols)
 
@@ -567,8 +572,8 @@ def _sql_as_cached_temp_table(
             job_config=job_config,
         )
         assert query_job is not None
-        query_job.result()
-        return query_job.destination
+        iter = query_job.result()
+        return query_job.destination, iter.total_rows
 
     def _validate_result_schema(
         self,

bigframes/version.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2.9.0"
+__version__ = "2.10.0"
 
 # {x-release-please-start-date}
-__release_date__ = "2025-06-30"
+__release_date__ = "2025-07-08"
 # {x-release-please-end}