cleanup

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -829,10 +829,12 @@ def _materialize_local(
                 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
-                return self._downsample(
-                    fraction=fraction,
+                down_sampled_block = self.split(
+                    fracs=(fraction,),
                     random_state=sample_config.random_state,
-                )._materialize_local(
+                    sort=False,
+                )[0]
+                return down_sampled_block._materialize_local(
                     MaterializationOptions(ordered=materialize_options.ordered)
                 )
         else:
@@ -841,14 +843,6 @@ def _materialize_local(
         df.set_axis(self.column_labels, axis=1, copy=False)
         return df, execute_result.query_job
 
-    def _downsample(self, fraction: float, random_state) -> Block:
-        block = self.split(
-            fracs=(fraction,),
-            random_state=random_state,
-            sort=False,
-        )[0]
-        return block
-
     def split(
         self,
         ns: Iterable[int] = (),

tests/system/small/test_anywidget.py

@@ -132,7 +132,7 @@ def execution_metadata(self) -> ExecutionMetadata:
         def schema(self):
             return schema
 
-        def batches(self) -> ResultsIterator:
+        def batches(self, sample_rate=None) -> ResultsIterator:
             return ResultsIterator(
                 arrow_batches_val,
                 self.schema,