Repository: ray-project/ray
File: python/ray/data/namespace_expressions/list_namespace.py (Python)

Code before:
    def flatten(self) -> "UDFExpr":
        """Flatten one level of nesting for each list value."""

        return_dtype = DataType(object)
        if self._expr.data_type.is_arrow_type():
            arrow_type = self._expr.data_type.to_arrow_dtype()
            if pyarrow.types.is_list(arrow_type) or pyarrow.types.is_large_list(
                arrow_type
            ):
                child_type = arrow_type.value_type
                list_factory = (
                    pyarrow.large_list
                    if pyarrow.types.is_large_list(arrow_type)
                    else pyarrow.list_
                )
                if (
                    pyarrow.types.is_list(child_type)
                    or pyarrow.types.is_large_list(child_type)
                    or pyarrow.types.is_fixed_size_list(child_type)
                ):
                    flattened_type = list_factory(child_type.value_type)
                    return_dtype = DataType.from_arrow(flattened_type)
            elif pyarrow.types.is_fixed_size_list(arrow_type):
                child_type = arrow_type.value_type
                if (
                    pyarrow.types.is_list(child_type)
                    or pyarrow.types.is_large_list(child_type)
                    or pyarrow.types.is_fixed_size_list(child_type)
                ):
                    flattened_type = pyarrow.list_(child_type.value_type)
                    return_dtype = DataType.from_arrow(flattened_type)

        if hasattr(pc, "list_flatten"):

            @pyarrow_udf(return_dtype=return_dtype)
            def _list_flatten_arrow(arr: pyarrow.Array) -> pyarrow.Array:
                if isinstance(arr, pyarrow.ChunkedArray):
                    arr = arr.combine_chunks()

                arr_type = arr.type
                arr_is_list = pyarrow.types.is_list(
                    arr_type
                ) or pyarrow.types.is_large_list(arr_type)
                arr_is_fixed_size = pyarrow.types.is_fixed_size_list(arr_type)
                if not (arr_is_list or arr_is_fixed_size):
                    raise TypeError(
                        "list.flatten() requires a list column whose elements are also lists."
                    )

                child_type = arr_type.value_type
                child_is_list = pyarrow.types.is_list(

Review comment (suggestion):
```
import pyarrow
import pyarrow.compute as pc

from ray.data.datatype import DataType


def _infer_flattened_dtype(expr: "Expr") -> DataType:
    """Infer the return DataType after flattening one level of list nesting."""
    if not expr.data_type.is_arrow_type():
        return DataType(object)

    arrow_type = expr.data_type.to_arrow_dtype()
    outer_dtype = DataType.from_arrow(arrow_type)
    if not outer_dtype.is_list_type():
        return DataType(object)

    child_type = arrow_type.value_type
    child_dtype = DataType.from_arrow(child_type)
    if not child_dtype.is_list_type():
        return DataType(object)

    if pyarrow.types.is_large_list(arrow_type):
        return DataType.from_arrow(pyarrow.large_list(child_type.value_type))
    else:
        return DataType.from_arrow(pyarrow.list_(child_type.value_type))


def _validate_nested_list(arr_type: pyarrow.DataType) -> None:
    """Raise TypeError if arr_type is not a list of lists."""
    outer_dtype = DataType.from_arrow(arr_type)
    if not outer_dtype.is_list_type():
        raise TypeError(
            "list.flatten() requires a list column whose elements are also lists."
        )

    child_dtype = DataType.from_arrow(arr_type.value_type)
    if not child_dtype.is_list_type():
        raise TypeError(
            "list.flatten() requires a list column whose elements are also lists."
        )


def flatten(self) -> "UDFExpr":
    """Flatten one level of nesting for each list value.

    Transforms list<list<T>> -> list<T> by concatenating inner lists per row.
    Example: [[1, 2], [3, 4]] -> [1, 2, 3, 4]
    """
    return_dtype = _infer_flattened_dtype(self._expr)

    @pyarrow_udf(return_dtype=return_dtype)
    def _list_flatten(arr: pyarrow.Array) -> pyarrow.Array:
        if isinstance(arr, pyarrow.ChunkedArray):
            arr = arr.combine_chunks()

        _validate_nested_list(arr.type)

        # Flatten: list<list<T>> -> all inner lists -> all scalar values
        inner_lists: pyarrow.Array = pc.list_flatten(arr)
        all_scalars: pyarrow.Array = pc.list_flatten(inner_lists)

        n_rows: int = len(arr)
        if len(all_scalars) == 0:
            offsets: pyarrow.Array = pyarrow.repeat(0, n_rows + 1)
        else:
            # Reconstruct row boundaries after flattening.
            #
            # pc.list_flatten loses row structure, so we must track which scalars
            # belong to which original row and rebuild the offsets.
            #
            # Example:
            #   arr = [[[1, 2], [3]], [[4, 5]]]  (2 rows)
            #   inner_lists = pc.list_flatten(arr) = [[1, 2], [3], [4, 5]]
            #   all_scalars = pc.list_flatten(inner_lists) = [1, 2, 3, 4, 5]
            #
            #   list_parent_indices(arr) = [0, 0, 1]  (inner list -> row)
            #   list_parent_indices(inner_lists) = [0, 0, 1, 2, 2]  (scalar -> inner list)
            #   row_indices = take([0, 0, 1], [0, 0, 1, 2, 2]) = [0, 0, 0, 1, 1]  (scalar -> row)
            #
            #   value_counts([0, 0, 0, 1, 1]) = {values: [0, 1], counts: [3, 2]}
            #   cumsum([3, 2]) = [3, 5]
            #   offsets = [0, 3, 5]  -> row 0 has scalars 0:3, row 1 has scalars 3:5
            #
            # Result: [[1, 2, 3], [4, 5]]

            # Map each scalar to its original row
            row_indices: pyarrow.Array = pc.take(
                pc.list_parent_indices(arr),
                pc.list_parent_indices(inner_lists),
            )

            # Count scalars per row using value_counts
            # Note: value_counts omits rows with 0 scalars, so we use index_in to handle gaps
            vc: pyarrow.StructArray = pc.value_counts(row_indices)
            rows_with_scalars: pyarrow.Array = pc.struct_field(vc, "values")
            scalar_counts: pyarrow.Array = pc.struct_field(vc, "counts")

            # For each row 0..n_rows-1, look up its count (0 if not present)
            row_sequence: pyarrow.Array = pyarrow.arange(0, n_rows)
            positions: pyarrow.Array = pc.index_in(row_sequence, value_set=rows_with_scalars)

            counts: pyarrow.Array = pc.if_else(
                pc.is_null(positions),
                0,
                pc.take(scalar_counts, pc.fill_null(positions, 0)),
            )

            # Build offsets: [0] + cumsum(counts)
            cumsum: pyarrow.Array = pc.cumulative_sum(counts)
            offsets = pc.concat_arrays([
                pyarrow.array([0], type=cumsum.type),
                cumsum,
            ])

        # Construct result with appropriate list type
        is_large: bool = pyarrow.types.is_large_list(arr.type)
        if not is_large:
            offsets = pc.cast(offsets, pyarrow.int32())

        null_mask: pyarrow.Array | None = arr.is_null() if arr.null_count else None
        array_cls: type = pyarrow.LargeListArray if is_large else pyarrow.ListArray
        return array_cls.from_arrays(offsets, all_scalars, mask=null_mask)

    return _list_flatten(self._expr)
```

Couple of things. 
1. Let's assume that `list_flatten` is a valid pyrrow function
2. Let's avoid the conversion to numpy from pyarrow
3. Let's leverage the existing functions that are in place in datatype.py 

Code after:

        return _list_slice(self._expr)

    def sort(
        self,
        order: Literal["ascending", "descending"] = "ascending",
        null_placement: Literal["at_start", "at_end"] = "at_end",
    ) -> "UDFExpr":
        """Sort the elements within each (nested) list.

        Args:
            order: Sorting order, must be ``\"ascending\"`` or ``\"descending\"``.
            null_placement: Placement for null values, ``\"at_start\"`` or ``\"at_end\"``.

        Returns:
            UDFExpr providing the sorted lists.

        Example:
            >>> from ray.data.expressions import col
            >>> # [[3,1],[2,None]] -> [[1,3],[2,None]]
            >>> expr = col("items").list.sort()  # doctest: +SKIP
        """

        if order not in {"ascending", "descending"}:
            raise ValueError(
                "order must be either 'ascending' or 'descending', got " f"{order!r}"
            )
        if null_placement not in {"at_start", "at_end"}:
            raise ValueError(
                "null_placement must be 'at_start' or 'at_end', got "
                f"{null_placement!r}"
            )

        return_dtype = self._expr.data_type

        @pyarrow_udf(return_dtype=return_dtype)
        def _list_sort(arr: pyarrow.Array) -> pyarrow.Array:
            # Approach:
            # 1) Normalize fixed_size_list -> list for list_* kernels (preserve nulls).
            # 2) Flatten to (row_index, value) pairs, sort by row then value.
            # 3) Rebuild list array using per-row lengths and restore original type.
            arr = _ensure_array(arr)

            arr_type = arr.type
            arr_dtype = DataType.from_arrow(arr_type)
            if not arr_dtype.is_list_type():
                raise TypeError("list.sort() requires a list column.")

            original_type = arr_type
            null_mask = arr.is_null() if arr.null_count else None
            sort_arr = arr