perf: Avoid requery for some result downsample methods (#2219)

Co-authored-by: Chelsea Lin <chelsealin@google.com>

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -818,49 +818,30 @@ def _materialize_local(
 total_rows=result_batches.approx_total_rows
 # Remove downsampling config from subsequent invocations, as otherwise could result in many
 # iterations if downsampling undershoots
-returnself._downsample(
-total_rows=total_rows,
-sampling_method=sample_config.sampling_method,
-fraction=fraction,
-random_state=sample_config.random_state,
- )._materialize_local(
-MaterializationOptions(ordered=materialize_options.ordered)
- )
-else:
-df=result_batches.to_pandas()
-df=self._copy_index_to_pandas(df)
-df.set_axis(self.column_labels, axis=1, copy=False)
-returndf, execute_result.query_job
-
-def_downsample(
-self, total_rows: int, sampling_method: str, fraction: float, random_state
- ) ->Block:
-# either selecting fraction or number of rows
-ifsampling_method==_HEAD:
-filtered_block=self.slice(stop=int(total_rows*fraction))
-returnfiltered_block
-elif (sampling_method==_UNIFORM) and (random_stateisNone):
-filtered_expr=self.expr._uniform_sampling(fraction)
-block=Block(
-filtered_expr,
-index_columns=self.index_columns,
-column_labels=self.column_labels,
-index_labels=self.index.names,
- )
-returnblock
-elifsampling_method==_UNIFORM:
-block=self.split(
-fracs=(fraction,),
-random_state=random_state,
-sort=False,
- )[0]
-returnblock
+ifsample_config.sampling_method=="head":
+# Just truncates the result iterator without a follow-up query
+raw_df=result_batches.to_pandas(limit=int(total_rows*fraction))
+elif (
+sample_config.sampling_method=="uniform"
+andsample_config.random_stateisNone
+ ):
+# Pushes sample into result without new query
+sampled_batches=execute_result.batches(sample_rate=fraction)
+raw_df=sampled_batches.to_pandas()
+else: # uniform sample with random state requires a full follow-up query
+down_sampled_block=self.split(
+fracs=(fraction,),
+random_state=sample_config.random_state,
+sort=False,
+ )[0]
+returndown_sampled_block._materialize_local(
+MaterializationOptions(ordered=materialize_options.ordered)
+ )
 else:
-# This part should never be called, just in case.
-raiseNotImplementedError(
-f"The downsampling method {sampling_method} is not implemented, "
-f"please choose from {','.join(_SAMPLING_METHODS)}."
- )
+raw_df=result_batches.to_pandas()
+df=self._copy_index_to_pandas(raw_df)
+df.set_axis(self.column_labels, axis=1, copy=False)
+returndf, execute_result.query_job
 
 defsplit(
 self,

bigframes/core/bq_data.py

@@ -186,11 +186,22 @@ def get_arrow_batches(
 columns: Sequence[str],
 storage_read_client: bigquery_storage_v1.BigQueryReadClient,
 project_id: str,
+sample_rate: Optional[float] =None,
 ) ->ReadResult:
 table_mod_options= {}
 read_options_dict: dict[str, Any] = {"selected_fields": list(columns)}
+
+predicates= []
 ifdata.sql_predicate:
-read_options_dict["row_restriction"] =data.sql_predicate
+predicates.append(data.sql_predicate)
+ifsample_rateisnotNone:
+assertisinstance(sample_rate, float)
+predicates.append(f"RAND() < {sample_rate}")
+
+ifpredicates:
+full_predicates=" AND ".join(f"( {pred} )"forpredinpredicates)
+read_options_dict["row_restriction"] =full_predicates
+
 read_options=bq_storage_types.ReadSession.TableReadOptions(**read_options_dict)
 
 ifdata.at_time:

bigframes/core/local_data.py

@@ -25,6 +25,7 @@
 importuuid
 
 importgeopandas# type: ignore
+importnumpy
 importnumpyasnp
 importpandasaspd
 importpyarrowaspa
@@ -124,13 +125,21 @@ def to_arrow(
 geo_format: Literal["wkb", "wkt"] ="wkt",
 duration_type: Literal["int", "duration"] ="duration",
 json_type: Literal["string"] ="string",
+sample_rate: Optional[float] =None,
 max_chunksize: Optional[int] =None,
  ) ->tuple[pa.Schema, Iterable[pa.RecordBatch]]:
 ifgeo_format!="wkt":
 raiseNotImplementedError(f"geo format {geo_format} not yet implemented")
 assertjson_type=="string"
 
-batches=self.data.to_batches(max_chunksize=max_chunksize)
+data=self.data
+
+# This exists for symmetry with remote sources, but sampling local data like this shouldn't really happen
+ifsample_rateisnotNone:
+to_take=numpy.random.rand(data.num_rows) <sample_rate
+data=data.filter(to_take)
+
+batches=data.to_batches(max_chunksize=max_chunksize)
 schema=self.data.schema
 ifduration_type=="int":
 schema=_schema_durations_to_ints(schema)

bigframes/session/executor.py

@@ -88,7 +88,7 @@ def arrow_batches(self) -> Iterator[pyarrow.RecordBatch]:
 
 yieldbatch
 
-defto_arrow_table(self) ->pyarrow.Table:
+defto_arrow_table(self, limit: Optional[int] =None) ->pyarrow.Table:
 # Need to provide schema if no result rows, as arrow can't infer
 # If ther are rows, it is safest to infer schema from batches.
 # Any discrepencies between predicted schema and actual schema will produce errors.
@@ -97,18 +97,21 @@ def to_arrow_table(self) -> pyarrow.Table:
 peek_value=list(peek_it)
 # TODO: Enforce our internal schema on the table for consistency
 iflen(peek_value) >0:
-returnpyarrow.Table.from_batches(
-itertools.chain(peek_value, batches), # reconstruct
- )
+batches=itertools.chain(peek_value, batches) # reconstruct
+iflimit:
+batches=pyarrow_utils.truncate_pyarrow_iterable(
+batches, max_results=limit
+ )
+returnpyarrow.Table.from_batches(batches)
 else:
 try:
 returnself._schema.to_pyarrow().empty_table()
 exceptpa.ArrowNotImplementedError:
 # Bug with some pyarrow versions, empty_table only supports base storage types, not extension types.
 returnself._schema.to_pyarrow(use_storage_types=True).empty_table()
 
-defto_pandas(self) ->pd.DataFrame:
-returnio_pandas.arrow_to_pandas(self.to_arrow_table(), self._schema)
+defto_pandas(self, limit: Optional[int] =None) ->pd.DataFrame:
+returnio_pandas.arrow_to_pandas(self.to_arrow_table(limit=limit), self._schema)
 
 defto_pandas_batches(
 self, page_size: Optional[int] =None, max_results: Optional[int] =None
@@ -158,7 +161,7 @@ def schema(self) -> bigframes.core.schema.ArraySchema:
  ...
 
 @abc.abstractmethod
-defbatches(self) ->ResultsIterator:
+defbatches(self, sample_rate: Optional[float] =None) ->ResultsIterator:
  ...
 
 @property
@@ -200,9 +203,9 @@ def execution_metadata(self) -> ExecutionMetadata:
 defschema(self) ->bigframes.core.schema.ArraySchema:
 returnself._data.schema
 
-defbatches(self) ->ResultsIterator:
+defbatches(self, sample_rate: Optional[float] =None) ->ResultsIterator:
 returnResultsIterator(
-iter(self._data.to_arrow()[1]),
+iter(self._data.to_arrow(sample_rate=sample_rate)[1]),
 self.schema,
 self._data.metadata.row_count,
 self._data.metadata.total_bytes,
@@ -226,7 +229,7 @@ def execution_metadata(self) -> ExecutionMetadata:
 defschema(self) ->bigframes.core.schema.ArraySchema:
 returnself._schema
 
-defbatches(self) ->ResultsIterator:
+defbatches(self, sample_rate: Optional[float] =None) ->ResultsIterator:
 returnResultsIterator(iter([]), self.schema, 0, 0)
 
 
@@ -260,12 +263,13 @@ def schema(self) -> bigframes.core.schema.ArraySchema:
 source_ids= [selection[0] forselectioninself._selected_fields]
 returnself._data.schema.select(source_ids).rename(dict(self._selected_fields))
 
-defbatches(self) ->ResultsIterator:
+defbatches(self, sample_rate: Optional[float] =None) ->ResultsIterator:
 read_batches=bq_data.get_arrow_batches(
 self._data,
  [x[0] forxinself._selected_fields],
 self._storage_client,
 self._project_id,
+sample_rate=sample_rate,
  )
 arrow_batches: Iterator[pa.RecordBatch] =map(
 functools.partial(

tests/system/small/test_anywidget.py

@@ -165,7 +165,7 @@ def execution_metadata(self) -> ExecutionMetadata:
 defschema(self) ->Any:
 returnschema
 
-defbatches(self) ->ResultsIterator:
+defbatches(self, sample_rate=None) ->ResultsIterator:
 returnResultsIterator(
 arrow_batches_val,
 self.schema,

tests/system/small/test_dataframe.py

@@ -4524,7 +4524,7 @@ def test_df_kurt(scalars_dfs):
 "n_default",
  ],
 )
-deftest_sample(scalars_dfs, frac, n, random_state):
+deftest_df_to_pandas_sample(scalars_dfs, frac, n, random_state):
 scalars_df, _=scalars_dfs
 df=scalars_df.sample(frac=frac, n=n, random_state=random_state)
 bf_result=df.to_pandas()
@@ -4535,15 +4535,15 @@ def test_sample(scalars_dfs, frac, n, random_state):
 assertbf_result.shape[1] ==scalars_df.shape[1]
 
 
-deftest_sample_determinism(penguins_df_default_index):
+deftest_df_to_pandas_sample_determinism(penguins_df_default_index):
 df=penguins_df_default_index.sample(n=100, random_state=12345).head(15)
 bf_result=df.to_pandas()
 bf_result2=df.to_pandas()
 
 pandas.testing.assert_frame_equal(bf_result, bf_result2)
 
 
-deftest_sample_raises_value_error(scalars_dfs):
+deftest_df_to_pandas_sample_raises_value_error(scalars_dfs):
 scalars_df, _=scalars_dfs
 withpytest.raises(
 ValueError, match="Only one of 'n' or 'frac' parameter can be specified."