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

Code before:
            )

        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 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.
        """

        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:
            if isinstance(arr, pyarrow.ChunkedArray):
                arr = arr.combine_chunks()

            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
            if pyarrow.types.is_fixed_size_list(arr_type):
                child_type = arr_type.value_type
                list_size = arr_type.list_size
                if null_mask is not None:
                    filler_values = pyarrow.nulls(len(arr) * list_size, type=child_type)
                    filler = pyarrow.FixedSizeListArray.from_arrays(
                        filler_values, list_size

Review comment (refactor):
One more comment:
```
def _ensure_array(arr: pyarrow.Array) -> pyarrow.Array:
    """Convert ChunkedArray to Array if needed."""
    if isinstance(arr, pyarrow.ChunkedArray):
        return arr.combine_chunks()
    return arr


def _build_list_array(
    offsets: pyarrow.Array,
    values: pyarrow.Array,
    is_large: bool,
    null_mask: pyarrow.Array | None = None,
) -> pyarrow.Array:
    """Reconstruct a ListArray from offsets and values."""
    offsets_type = pyarrow.int64() if is_large else pyarrow.int32()
    offsets = pc.cast(offsets, offsets_type)
    array_cls = pyarrow.LargeListArray if is_large else pyarrow.ListArray
    return array_cls.from_arrays(offsets, values, mask=null_mask)


def _counts_to_offsets(counts: pyarrow.Array) -> pyarrow.Array:
    """Convert per-row counts to list offsets via cumulative sum."""
    cumsum = pc.cumulative_sum(counts)
    return pyarrow.concat_arrays([pyarrow.array([0], type=cumsum.type), cumsum])
```

I realized sort and flatten are handling chunkedarrays, list reconstruction, and offset computation very similarly. With these helpers, the 2 functions can be simplified. 

Code after:
    ) -> "UDFExpr":
        """Slice each list.

        Args:
            start: Start index (inclusive). Defaults to 0.
            stop: Stop index (exclusive). Defaults to list length.
            step: Step size. Defaults to 1.

        Returns:
            UDFExpr that extracts a slice from each list.
        """
        # Return type is the same as the input list type
        return_dtype = self._expr.data_type

        @pyarrow_udf(return_dtype=return_dtype)
        def _list_slice(arr: pyarrow.Array) -> pyarrow.Array:
            return pc.list_slice(
                arr,
                start=0 if start is None else start,
                stop=stop,
                step=1 if step is None else step,
            )

        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(