wehub-resource-sync
336 lines
13 KiB
Python
336 lines
13 KiB
Python
"""List namespace for expression operations on list-typed columns."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Literal, Union
|
|
|
|
import numpy as np
|
|
import pyarrow
|
|
import pyarrow.compute as pc
|
|
|
|
from ray.data._internal.arrow_utils import _combine_as_list_array, _counts_to_offsets
|
|
from ray.data.datatype import DataType
|
|
from ray.data.expressions import _create_pyarrow_compute_udf, pyarrow_udf
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.expressions import Expr, PyArrowComputeUDFExpr, UDFExpr
|
|
|
|
|
|
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 _is_list_like(pa_type: pyarrow.DataType) -> bool:
|
|
"""Return True for list-like Arrow types (list, large_list, fixed_size_list)."""
|
|
return (
|
|
pyarrow.types.is_list(pa_type)
|
|
or pyarrow.types.is_large_list(pa_type)
|
|
or pyarrow.types.is_fixed_size_list(pa_type)
|
|
or (
|
|
hasattr(pyarrow.types, "is_list_view")
|
|
and pyarrow.types.is_list_view(pa_type)
|
|
)
|
|
or (
|
|
hasattr(pyarrow.types, "is_large_list_view")
|
|
and pyarrow.types.is_large_list_view(pa_type)
|
|
)
|
|
)
|
|
|
|
|
|
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()
|
|
if not _is_list_like(arrow_type):
|
|
return DataType(object)
|
|
|
|
child_type = arrow_type.value_type
|
|
if not _is_list_like(child_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."""
|
|
if not _is_list_like(arr_type):
|
|
raise TypeError(
|
|
"list.flatten() requires a list column whose elements are also lists."
|
|
)
|
|
|
|
if not _is_list_like(arr_type.value_type):
|
|
raise TypeError(
|
|
"list.flatten() requires a list column whose elements are also lists."
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _ListNamespace:
|
|
"""Namespace for list operations on expression columns.
|
|
|
|
This namespace provides methods for operating on list-typed columns using
|
|
PyArrow compute functions.
|
|
|
|
Example:
|
|
>>> from ray.data.expressions import col
|
|
>>> # Get length of list column
|
|
>>> expr = col("items").list.len()
|
|
>>> # Get first item using method
|
|
>>> expr = col("items").list.get(0)
|
|
>>> # Get first item using indexing
|
|
>>> expr = col("items").list[0]
|
|
>>> # Slice list
|
|
>>> expr = col("items").list[1:3]
|
|
"""
|
|
|
|
_expr: Expr
|
|
|
|
def len(self) -> "PyArrowComputeUDFExpr":
|
|
"""Get the length of each list."""
|
|
return _create_pyarrow_compute_udf(pc.list_value_length, DataType.int32())(
|
|
self._expr
|
|
)
|
|
|
|
def __getitem__(self, key: Union[int, slice]) -> "PyArrowComputeUDFExpr":
|
|
"""Get element or slice using bracket notation.
|
|
|
|
Args:
|
|
key: An integer for element access or slice for list slicing.
|
|
|
|
Returns:
|
|
Expression that extracts the element or slice.
|
|
|
|
Example:
|
|
>>> col("items").list[0] # Get first item # doctest: +SKIP
|
|
>>> col("items").list[1:3] # Get slice [1, 3) # doctest: +SKIP
|
|
>>> col("items").list[-1] # Get last item # doctest: +SKIP
|
|
"""
|
|
if isinstance(key, int):
|
|
return self.get(key)
|
|
elif isinstance(key, slice):
|
|
return self.slice(key.start, key.stop, key.step)
|
|
else:
|
|
raise TypeError(
|
|
f"List indices must be integers or slices, not {type(key).__name__}"
|
|
)
|
|
|
|
def get(self, index: int) -> "PyArrowComputeUDFExpr":
|
|
"""Get element at the specified index from each list.
|
|
|
|
Args:
|
|
index: The index of the element to retrieve. Negative indices are supported.
|
|
|
|
Returns:
|
|
Expression that extracts the element at the given index.
|
|
"""
|
|
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)
|
|
or pyarrow.types.is_fixed_size_list(arrow_type)
|
|
):
|
|
return_dtype = DataType.from_arrow(arrow_type.value_type)
|
|
|
|
return _create_pyarrow_compute_udf(pc.list_element, return_dtype)(
|
|
self._expr, index
|
|
)
|
|
|
|
def slice(
|
|
self, start: int | None = None, stop: int | None = None, step: int | None = None
|
|
) -> "PyArrowComputeUDFExpr":
|
|
"""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:
|
|
Expression that extracts a slice from each list.
|
|
"""
|
|
return _create_pyarrow_compute_udf(pc.list_slice, self._expr.data_type)(
|
|
self._expr,
|
|
start=0 if start is None else start,
|
|
stop=stop,
|
|
step=1 if step is None else step,
|
|
)
|
|
|
|
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
|
|
if pyarrow.types.is_fixed_size_list(arr_type):
|
|
# Example: FixedSizeList<2>[ [3,1], None, [2,4] ]
|
|
# Fill null row -> [[3,1],[None,None],[2,4]], cast to list<child> for sort,
|
|
# then cast back to fixed_size to preserve schema. list_* kernels operate
|
|
# on list/large_list, so we cast fixed_size_list<T> to list<T> here.
|
|
child_type = arr_type.value_type
|
|
list_size = arr_type.list_size
|
|
if null_mask is not None:
|
|
# Fill null rows with fixed-size null lists so each row keeps
|
|
# the same list_size when we sort and rebuild offsets.
|
|
filler_values = pyarrow.nulls(len(arr) * list_size, type=child_type)
|
|
filler = pyarrow.FixedSizeListArray.from_arrays(
|
|
filler_values, list_size
|
|
)
|
|
sort_arr = pc.if_else(null_mask, filler, arr)
|
|
list_type = pyarrow.list_(child_type)
|
|
sort_arr = sort_arr.cast(list_type)
|
|
arr_type = sort_arr.type
|
|
# Flatten to (row_index, value) pairs, sort within each row by value.
|
|
values = pc.list_flatten(sort_arr)
|
|
if len(values):
|
|
row_indices = pc.list_parent_indices(sort_arr)
|
|
struct = pyarrow.StructArray.from_arrays(
|
|
[row_indices, values],
|
|
["row", "value"],
|
|
)
|
|
sorted_indices = pc.sort_indices(
|
|
struct,
|
|
sort_keys=[("row", "ascending"), ("value", order)],
|
|
null_placement=null_placement,
|
|
)
|
|
values = pc.take(values, sorted_indices)
|
|
|
|
# Reconstruct list array with original row boundaries and nulls.
|
|
lengths = pc.list_value_length(sort_arr)
|
|
lengths = pc.fill_null(lengths, 0)
|
|
is_large = pyarrow.types.is_large_list(arr_type)
|
|
offsets = _counts_to_offsets(lengths)
|
|
sorted_arr = _combine_as_list_array(
|
|
offsets=offsets,
|
|
values=values,
|
|
is_large=is_large,
|
|
null_mask=null_mask,
|
|
)
|
|
|
|
if pyarrow.types.is_fixed_size_list(original_type):
|
|
sorted_arr = sorted_arr.cast(original_type)
|
|
|
|
return sorted_arr
|
|
|
|
return _list_sort(self._expr)
|
|
|
|
def flatten(self) -> "UDFExpr":
|
|
"""Flatten one level of nesting for each list value."""
|
|
|
|
return_dtype = _infer_flattened_dtype(self._expr)
|
|
|
|
@pyarrow_udf(return_dtype=return_dtype)
|
|
def _list_flatten(arr: pyarrow.Array) -> pyarrow.Array:
|
|
# Approach:
|
|
# 1) Flatten list<list<T>> to a flat values array and parent indices.
|
|
# 2) Count values per original row.
|
|
# 3) Rebuild list array using offsets while preserving top-level nulls.
|
|
arr = _ensure_array(arr)
|
|
|
|
_validate_nested_list(arr.type)
|
|
|
|
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:
|
|
# All rows are empty/None after flatten, so build zero counts to
|
|
# preserve row count and produce empty lists for each row.
|
|
counts = pyarrow.array(np.repeat(0, n_rows), type=pyarrow.int64())
|
|
offsets = _counts_to_offsets(counts)
|
|
else:
|
|
# Example: arr = [[[1,2],[3]], [[4], None], None]
|
|
# inner_lists = [[1,2],[3],[4],None], all_scalars = [1,2,3,4]
|
|
# parent(arr)=[0,0,1,1], parent(inner)=[0,0,1,2] -> row_indices=[0,0,0,1]
|
|
# counts=[3,1,0] -> offsets=[0,3,4,4]
|
|
row_indices: pyarrow.Array = pc.take(
|
|
pc.list_parent_indices(arr),
|
|
pc.list_parent_indices(inner_lists),
|
|
)
|
|
|
|
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")
|
|
|
|
# Compute per-row counts of flattened scalars. value_counts gives counts
|
|
# only for rows that appear, so we map those counts back onto the full
|
|
# row range [0, n_rows) and fill missing rows with 0.
|
|
row_sequence: pyarrow.Array = pyarrow.array(
|
|
np.arange(n_rows, dtype=np.int64), type=pyarrow.int64()
|
|
)
|
|
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)),
|
|
)
|
|
|
|
offsets = _counts_to_offsets(counts)
|
|
|
|
is_large: bool = pyarrow.types.is_large_list(arr.type)
|
|
null_mask: pyarrow.Array | None = arr.is_null() if arr.null_count else None
|
|
# Rebuild a list/large_list array while preserving top-level nulls.
|
|
return _combine_as_list_array(
|
|
offsets=offsets,
|
|
values=all_scalars,
|
|
is_large=is_large,
|
|
null_mask=null_mask,
|
|
)
|
|
|
|
return _list_flatten(self._expr)
|