[Data] Abstracted BlockColumnAccessor (#51326)

## Why are these changes needed?

This component is generalizing and abstracting common columnar
operations.

Additionally 

 - Added tests
 - Cleaned up some minor issues,
 - Enable `concat` and `concat_and_sort` to do type promotions

---------

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Author: alexeykudinkin

python/ray/data/_internal/arrow_block.py

@@ -33,6 +33,7 @@
     BlockMetadata,
     BlockType,
     U,
+    BlockColumnAccessor,
 )
 from ray.data.context import DataContext
 
@@ -329,87 +330,6 @@ def _sample(self, n_samples: int, sort_key: "SortKey") -> "pyarrow.Table":
         table = self._table.select(sort_key.get_columns())
         return transform_pyarrow.take_table(table, indices)
 
-    def count(self, on: str, ignore_nulls: bool = False) -> Optional[U]:
-        """Count the number of non-null values in the provided column."""
-        import pyarrow.compute as pac
-
-        if not isinstance(on, str):
-            raise ValueError(
-                "on must be a string when aggregating on Arrow blocks, but got:"
-                f"{type(on)}."
-            )
-
-        if self.num_rows() == 0:
-            return None
-
-        mode = "only_valid" if ignore_nulls else "all"
-
-        col = self._table[on]
-        return pac.count(col, mode=mode).as_py()
-
-    def _apply_arrow_compute(
-        self, compute_fn: Callable, on: str, ignore_nulls: bool
-    ) -> Optional[U]:
-        """Helper providing null handling around applying an aggregation to a column."""
-        import pyarrow as pa
-
-        if not isinstance(on, str):
-            raise ValueError(
-                "on must be a string when aggregating on Arrow blocks, but got:"
-                f"{type(on)}."
-            )
-
-        if self.num_rows() == 0:
-            return None
-
-        col = self._table[on]
-        if pa.types.is_null(col.type):
-            return None
-        else:
-            return compute_fn(col, skip_nulls=ignore_nulls).as_py()
-
-    def sum(self, on: str, ignore_nulls: bool) -> Optional[U]:
-        import pyarrow.compute as pac
-
-        return self._apply_arrow_compute(pac.sum, on, ignore_nulls)
-
-    def min(self, on: str, ignore_nulls: bool) -> Optional[U]:
-        import pyarrow.compute as pac
-
-        return self._apply_arrow_compute(pac.min, on, ignore_nulls)
-
-    def max(self, on: str, ignore_nulls: bool) -> Optional[U]:
-        import pyarrow.compute as pac
-
-        return self._apply_arrow_compute(pac.max, on, ignore_nulls)
-
-    def mean(self, on: str, ignore_nulls: bool) -> Optional[U]:
-        import pyarrow.compute as pac
-
-        return self._apply_arrow_compute(pac.mean, on, ignore_nulls)
-
-    def sum_of_squared_diffs_from_mean(
-        self,
-        on: str,
-        ignore_nulls: bool,
-        mean: Optional[U] = None,
-    ) -> Optional[U]:
-        import pyarrow.compute as pac
-
-        if mean is None:
-            # If precomputed mean not given, we compute it ourselves.
-            mean = self.mean(on, ignore_nulls)
-            if mean is None:
-                return None
-        return self._apply_arrow_compute(
-            lambda col, skip_nulls: pac.sum(
-                pac.power(pac.subtract(col, mean), 2),
-                skip_nulls=skip_nulls,
-            ),
-            on,
-            ignore_nulls,
-        )
-
     def sort(self, sort_key: "SortKey") -> Block:
         assert (
             sort_key.get_columns()
@@ -449,8 +369,63 @@ def merge_sorted_blocks(
             # Handle blocks of different types.
             blocks = TableBlockAccessor.normalize_block_types(blocks, BlockType.ARROW)
             concat_and_sort = get_concat_and_sort_transform(DataContext.get_current())
-            ret = concat_and_sort(blocks, sort_key)
+            ret = concat_and_sort(blocks, sort_key, promote_types=True)
         return ret, ArrowBlockAccessor(ret).get_metadata(exec_stats=stats.build())
 
     def block_type(self) -> BlockType:
         return BlockType.ARROW
+
+
+class ArrowBlockColumnAccessor(BlockColumnAccessor):
+    def __init__(self, col: Union["pyarrow.Array", "pyarrow.ChunkedArray"]):
+        super().__init__(col)
+
+    def count(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.count(self._column, mode="only_valid" if ignore_nulls else "all")
+        return res.as_py() if as_py else res
+
+    def sum(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.sum(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def min(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.min(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def max(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.max(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def mean(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.mean(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def sum_of_squared_diffs_from_mean(
+        self, ignore_nulls: bool, mean: Optional[U] = None, as_py: bool = True
+    ) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        # Calculate mean if not provided
+        if mean is None:
+            mean = self.mean(ignore_nulls=ignore_nulls)
+
+        if mean is None:
+            return None
+
+        res = pac.sum(
+            pac.power(pac.subtract(self._column, mean), 2), skip_nulls=ignore_nulls
+        )
+        return res.as_py() if as_py else res
+
+    def to_pylist(self):
+        return self._column.to_pylist()

python/ray/data/_internal/arrow_ops/transform_polars.py

@@ -30,7 +30,7 @@ def sort(table: "pyarrow.Table", sort_key: "SortKey") -> "pyarrow.Table":
 
 
 def concat_and_sort(
-    blocks: List["pyarrow.Table"], sort_key: "SortKey"
+    blocks: List["pyarrow.Table"], sort_key: "SortKey", *, promote_types: bool = False
 ) -> "pyarrow.Table":
     check_polars_installed()
     blocks = [pl.from_arrow(block) for block in blocks]

python/ray/data/_internal/arrow_ops/transform_pyarrow.py

@@ -454,7 +454,7 @@ def concat(
 
     if not blocks:
         # Short-circuit on empty list of blocks.
-        return blocks
+        return pa.table([])
 
     if len(blocks) == 1:
         return blocks[0]
@@ -566,12 +566,20 @@ def concat(
 
 
 def concat_and_sort(
-    blocks: List["pyarrow.Table"], sort_key: "SortKey"
+    blocks: List["pyarrow.Table"],
+    sort_key: "SortKey",
+    *,
+    promote_types: bool = False,
 ) -> "pyarrow.Table":
+    import pyarrow as pa
     import pyarrow.compute as pac
 
-    ret = concat(blocks, promote_types=True)
+    if len(blocks) == 0:
+        return pa.table([])
+
+    ret = concat(blocks, promote_types=promote_types)
     indices = pac.sort_indices(ret, sort_keys=sort_key.to_arrow_sort_args())
+
     return take_table(ret, indices)