[Data] Compute Expressions-fixed size array (ray-project#58741)

## Description
Completing the fixed-size array namespace operations

## Related issues
Related to ray-project#58674

## Additional information

---------

Signed-off-by: 400Ping <fourhundredping@gmail.com>
Signed-off-by: Ping <fourhundredping@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: peterxcli <peterxcli@gmail.com>

Author: 2 people

python/ray/data/expressions.py

@@ -26,6 +26,7 @@
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
 if TYPE_CHECKING:
+    from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace
     from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace
     from ray.data.namespace_expressions.list_namespace import _ListNamespace
     from ray.data.namespace_expressions.string_namespace import _StringNamespace
@@ -577,12 +578,20 @@ def abs(self) -> "UDFExpr":
         """
         return _create_pyarrow_compute_udf(pc.abs_checked)(self)
 
+    @property
+    def arr(self) -> "_ArrayNamespace":
+        """Access array operations for this expression."""
+        from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace
+
+        return _ArrayNamespace(self)
+
     @property
     def list(self) -> "_ListNamespace":
         """Access list operations for this expression.
 
         Returns:
-            A _ListNamespace that provides list-specific operations.
+            A _ListNamespace that provides list-specific operations for both
+            PyArrow ``List`` and ``FixedSizeList`` columns.
 
         Example:
             >>> from ray.data.expressions import col
@@ -1490,6 +1499,7 @@ def download(uri_column_name: str) -> DownloadExpr:
     "lit",
     "download",
     "star",
+    "_ArrayNamespace",
     "_ListNamespace",
     "_StringNamespace",
     "_StructNamespace",
@@ -1499,7 +1509,11 @@ def download(uri_column_name: str) -> DownloadExpr:
 
 def __getattr__(name: str):
     """Lazy import of namespace classes to avoid circular imports."""
-    if name == "_ListNamespace":
+    if name == "_ArrayNamespace":
+        from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace
+
+        return _ArrayNamespace
+    elif name == "_ListNamespace":
         from ray.data.namespace_expressions.list_namespace import _ListNamespace
 
         return _ListNamespace

python/ray/data/namespace_expressions/arr_namespace.py

@@ -0,0 +1,54 @@
+"""Array namespace for expression operations on array-typed columns."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import pyarrow
+
+from ray.data.datatype import DataType
+from ray.data.expressions import pyarrow_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, UDFExpr
+
+
+@dataclass
+class _ArrayNamespace:
+    """Namespace for array operations on expression columns.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> # Convert fixed-size lists to variable-length lists
+        >>> expr = col("features").arr.to_list()
+    """
+
+    _expr: Expr
+
+    def to_list(self) -> "UDFExpr":
+        """Convert FixedSizeList columns into variable-length lists."""
+        return_dtype = DataType(object)
+
+        expr_dtype = self._expr.data_type
+        if expr_dtype.is_list_type():
+            arrow_type = expr_dtype.to_arrow_dtype()
+            if pyarrow.types.is_fixed_size_list(arrow_type):
+                return_dtype = DataType.from_arrow(pyarrow.list_(arrow_type.value_type))
+            else:
+                return_dtype = expr_dtype
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _to_list(arr: pyarrow.Array) -> pyarrow.Array:
+            arr_dtype = DataType.from_arrow(arr.type)
+            if not arr_dtype.is_list_type():
+                raise pyarrow.lib.ArrowInvalid(
+                    "to_list() can only be called on list-like columns, "
+                    f"but got {arr.type}"
+                )
+
+            if isinstance(arr.type, pyarrow.FixedSizeListType):
+                return arr.cast(pyarrow.list_(arr.type.value_type))
+            return arr
+
+        return _to_list(self._expr)

python/ray/data/tests/expressions/test_namespace_arr.py

@@ -0,0 +1,63 @@
+"""Integration tests for array namespace expressions.
+
+These tests require Ray and test end-to-end array namespace expression evaluation.
+"""
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+from packaging import version
+
+import ray
+from ray.data._internal.util import rows_same
+from ray.data.expressions import col
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+pytestmark = pytest.mark.skipif(
+    version.parse(pa.__version__) < version.parse("19.0.0"),
+    reason="Namespace expressions tests require PyArrow >= 19.0",
+)
+
+
+def _make_fixed_size_list_table() -> pa.Table:
+    values = pa.array([1, 2, 3, 4, 5, 6], type=pa.int64())
+    fixed = pa.FixedSizeListArray.from_arrays(values, list_size=2)
+    return pa.Table.from_arrays([fixed], names=["features"])
+
+
+def test_arr_to_list_fixed_size(ray_start_regular_shared):
+    table = _make_fixed_size_list_table()
+    ds = ray.data.from_arrow(table)
+
+    result = (
+        ds.with_column("features", col("features").arr.to_list())
+        .select_columns(["features"])
+        .to_pandas()
+    )
+    expected = pd.DataFrame(
+        [
+            {"features": [1, 2]},
+            {"features": [3, 4]},
+            {"features": [5, 6]},
+        ]
+    )
+
+    assert rows_same(result, expected)
+
+
+def test_arr_to_list_invalid_dtype_raises(ray_start_regular_shared):
+    ds = ray.data.from_items([{"value": 1}, {"value": 2}])
+
+    with pytest.raises(
+        (ray.exceptions.RayTaskError, ray.exceptions.UserCodeException)
+    ) as exc_info:
+        ds.with_column("value_list", col("value").arr.to_list()).to_pandas()
+
+    assert "to_list() can only be called on list-like columns" in str(exc_info.value)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))