chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
@@ -0,0 +1,87 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.expressions import _create_pyarrow_compute_udf
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.expressions import Expr, PyArrowComputeUDFExpr
|
||||
|
||||
TemporalUnit = Literal[
|
||||
"year",
|
||||
"quarter",
|
||||
"month",
|
||||
"week",
|
||||
"day",
|
||||
"hour",
|
||||
"minute",
|
||||
"second",
|
||||
"millisecond",
|
||||
"microsecond",
|
||||
"nanosecond",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DatetimeNamespace:
|
||||
"""Datetime namespace for operations on datetime-typed expression columns."""
|
||||
|
||||
_expr: "Expr"
|
||||
|
||||
# extractors
|
||||
|
||||
def year(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract year component."""
|
||||
return _create_pyarrow_compute_udf(pc.year, DataType.int32())(self._expr)
|
||||
|
||||
def month(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract month component."""
|
||||
return _create_pyarrow_compute_udf(pc.month, DataType.int32())(self._expr)
|
||||
|
||||
def day(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract day component."""
|
||||
return _create_pyarrow_compute_udf(pc.day, DataType.int32())(self._expr)
|
||||
|
||||
def hour(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract hour component."""
|
||||
return _create_pyarrow_compute_udf(pc.hour, DataType.int32())(self._expr)
|
||||
|
||||
def minute(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract minute component."""
|
||||
return _create_pyarrow_compute_udf(pc.minute, DataType.int32())(self._expr)
|
||||
|
||||
def second(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract second component."""
|
||||
return _create_pyarrow_compute_udf(pc.second, DataType.int32())(self._expr)
|
||||
|
||||
# formatting
|
||||
|
||||
def strftime(self, fmt: str) -> "PyArrowComputeUDFExpr":
|
||||
"""Format timestamps with a strftime pattern."""
|
||||
return _create_pyarrow_compute_udf(pc.strftime, DataType.string())(
|
||||
self._expr, format=fmt
|
||||
)
|
||||
|
||||
# rounding
|
||||
|
||||
def ceil(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
|
||||
"""Ceil timestamps to the next multiple of the given unit."""
|
||||
return _create_pyarrow_compute_udf(pc.ceil_temporal, self._expr.data_type)(
|
||||
self._expr, multiple=1, unit=unit
|
||||
)
|
||||
|
||||
def floor(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
|
||||
"""Floor timestamps to the previous multiple of the given unit."""
|
||||
return _create_pyarrow_compute_udf(pc.floor_temporal, self._expr.data_type)(
|
||||
self._expr, multiple=1, unit=unit
|
||||
)
|
||||
|
||||
def round(self, unit: TemporalUnit) -> "PyArrowComputeUDFExpr":
|
||||
"""Round timestamps to the nearest multiple of the given unit."""
|
||||
return _create_pyarrow_compute_udf(pc.round_temporal, self._expr.data_type)(
|
||||
self._expr, multiple=1, unit=unit
|
||||
)
|
||||
@@ -0,0 +1,335 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,193 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import numpy as np
|
||||
import pyarrow
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.expressions import pyarrow_udf
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.expressions import Expr, UDFExpr
|
||||
|
||||
|
||||
class MapComponent(str, Enum):
|
||||
KEYS = "keys"
|
||||
VALUES = "values"
|
||||
|
||||
|
||||
def _get_child_array(
|
||||
arr: pyarrow.Array, component: MapComponent
|
||||
) -> Optional[pyarrow.Array]:
|
||||
"""Extract the flat keys or values array from a map-like array.
|
||||
|
||||
Example: MapArray [{"a": 1}, {"b": 2}] -> keys ["a", "b"] or values [1, 2]
|
||||
"""
|
||||
if isinstance(arr, pyarrow.MapArray):
|
||||
if component == MapComponent.KEYS:
|
||||
return arr.keys
|
||||
else:
|
||||
return arr.items
|
||||
|
||||
if isinstance(arr, (pyarrow.ListArray, pyarrow.LargeListArray)):
|
||||
flat_values = arr.values
|
||||
if (
|
||||
isinstance(flat_values, pyarrow.StructArray)
|
||||
and flat_values.type.num_fields >= 2
|
||||
):
|
||||
idx = 0 if component == MapComponent.KEYS else 1
|
||||
return flat_values.field(idx)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _make_empty_list_array(
|
||||
arr: pyarrow.Array, component: MapComponent
|
||||
) -> pyarrow.Array:
|
||||
"""Create an all-null ListArray matching the input length.
|
||||
|
||||
Example: arr of length 3 -> ListArray [null, null, null]
|
||||
"""
|
||||
if len(arr) > 0 and arr.null_count < len(arr):
|
||||
raise TypeError(
|
||||
f"Expression is not a valid map type. .map.{component.value}() requires "
|
||||
f"pyarrow.MapArray or pyarrow.ListArray<Struct> with at least 2 fields "
|
||||
f"(key and value), but got: {arr.type}."
|
||||
)
|
||||
return pyarrow.ListArray.from_arrays(
|
||||
offsets=np.repeat(0, len(arr) + 1),
|
||||
values=pyarrow.array([], type=pyarrow.null()),
|
||||
mask=pyarrow.array(np.repeat(True, len(arr))),
|
||||
)
|
||||
|
||||
|
||||
def _rebuild_list_array(
|
||||
arr: pyarrow.Array, child_array: pyarrow.Array
|
||||
) -> pyarrow.Array:
|
||||
"""Rebuild a ListArray from parent offsets and child values, normalizing sliced offsets.
|
||||
|
||||
Example: offsets [5, 7, 10] -> slice child to [5:10], normalize offsets to [0, 2, 5]
|
||||
"""
|
||||
offsets = arr.offsets
|
||||
if len(offsets) > 0:
|
||||
start_offset = offsets[0]
|
||||
if start_offset != 0:
|
||||
end_offset = offsets[-1]
|
||||
child_array = child_array.slice(
|
||||
offset=int(start_offset), length=int(end_offset) - int(start_offset)
|
||||
)
|
||||
offsets = pc.subtract(offsets, start_offset)
|
||||
|
||||
factory = (
|
||||
pyarrow.LargeListArray.from_arrays
|
||||
if isinstance(arr, pyarrow.LargeListArray)
|
||||
else pyarrow.ListArray.from_arrays
|
||||
)
|
||||
return factory(offsets=offsets, values=child_array, mask=arr.is_null())
|
||||
|
||||
|
||||
def _get_result_type(
|
||||
arr_type: pyarrow.DataType, component: MapComponent
|
||||
) -> pyarrow.DataType:
|
||||
"""Infer the result list type from the input map type."""
|
||||
if pyarrow.types.is_map(arr_type):
|
||||
inner = (
|
||||
arr_type.key_type if component == MapComponent.KEYS else arr_type.item_type
|
||||
)
|
||||
return pyarrow.list_(inner)
|
||||
if pyarrow.types.is_list(arr_type) or pyarrow.types.is_large_list(arr_type):
|
||||
struct_type = arr_type.value_type
|
||||
if pyarrow.types.is_struct(struct_type) and struct_type.num_fields >= 2:
|
||||
idx = 0 if component == MapComponent.KEYS else 1
|
||||
list_fn = (
|
||||
pyarrow.large_list
|
||||
if pyarrow.types.is_large_list(arr_type)
|
||||
else pyarrow.list_
|
||||
)
|
||||
return list_fn(struct_type.field(idx).type)
|
||||
return pyarrow.list_(pyarrow.null())
|
||||
|
||||
|
||||
def _extract_map_component(
|
||||
arr: pyarrow.Array, component: MapComponent
|
||||
) -> pyarrow.Array:
|
||||
"""Extract keys or values from a MapArray or ListArray<Struct>.
|
||||
|
||||
This serves as the primary implementation since PyArrow does not yet
|
||||
expose dedicated compute kernels for map projection in the Python API.
|
||||
"""
|
||||
if isinstance(arr, pyarrow.ChunkedArray):
|
||||
chunks = [_extract_map_component(chunk, component) for chunk in arr.chunks]
|
||||
if not chunks:
|
||||
return pyarrow.chunked_array([], type=_get_result_type(arr.type, component))
|
||||
return pyarrow.chunked_array(chunks)
|
||||
|
||||
child_array = _get_child_array(arr, component)
|
||||
|
||||
if child_array is None:
|
||||
return _make_empty_list_array(arr, component)
|
||||
|
||||
return _rebuild_list_array(arr, child_array)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MapNamespace:
|
||||
"""Namespace for map operations on expression columns.
|
||||
|
||||
This namespace provides methods for operating on map-typed columns
|
||||
(including MapArrays and ListArrays of Structs) using PyArrow UDFs.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> # Get keys from map column
|
||||
>>> expr = col("headers").map.keys()
|
||||
>>> # Get values from map column
|
||||
>>> expr = col("headers").map.values()
|
||||
"""
|
||||
|
||||
_expr: "Expr"
|
||||
|
||||
def keys(self) -> "UDFExpr":
|
||||
"""Returns a list expression containing the keys of the map.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> # Get keys from map column
|
||||
>>> expr = col("headers").map.keys()
|
||||
|
||||
Returns:
|
||||
A list expression containing the keys.
|
||||
"""
|
||||
return self._create_projection_udf(MapComponent.KEYS)
|
||||
|
||||
def values(self) -> "UDFExpr":
|
||||
"""Returns a list expression containing the values of the map.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> # Get values from map column
|
||||
>>> expr = col("headers").map.values()
|
||||
|
||||
Returns:
|
||||
A list expression containing the values.
|
||||
"""
|
||||
return self._create_projection_udf(MapComponent.VALUES)
|
||||
|
||||
def _create_projection_udf(self, component: MapComponent) -> "UDFExpr":
|
||||
"""Helper to generate UDFs for map projections."""
|
||||
|
||||
return_dtype = DataType(object)
|
||||
if self._expr.data_type.is_arrow_type():
|
||||
arrow_type = self._expr.data_type.to_arrow_dtype()
|
||||
result_arrow_type = _get_result_type(arrow_type, component)
|
||||
return_dtype = DataType.from_arrow(result_arrow_type)
|
||||
|
||||
@pyarrow_udf(return_dtype=return_dtype)
|
||||
def _project_map(arr: pyarrow.Array) -> pyarrow.Array:
|
||||
return _extract_map_component(arr, component)
|
||||
|
||||
return _project_map(self._expr)
|
||||
@@ -0,0 +1,376 @@
|
||||
"""String namespace for expression operations on string-typed columns."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Callable, Literal
|
||||
|
||||
import pyarrow
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.expressions import _create_pyarrow_compute_udf
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.expressions import Expr, PyArrowComputeUDFExpr
|
||||
|
||||
|
||||
def _create_str_udf(
|
||||
pc_func: Callable[..., pyarrow.Array], return_dtype: DataType
|
||||
) -> Callable[..., "PyArrowComputeUDFExpr"]:
|
||||
"""Helper to create a string UDF that wraps a PyArrow compute function.
|
||||
|
||||
This helper handles all types of PyArrow compute operations:
|
||||
- Unary operations (no args): upper(), lower(), reverse()
|
||||
- Pattern operations (pattern + args): starts_with(), contains()
|
||||
- Multi-argument operations: replace(), replace_slice()
|
||||
|
||||
Args:
|
||||
pc_func: PyArrow compute function that takes (array, *positional, **kwargs)
|
||||
return_dtype: The return data type
|
||||
|
||||
Returns:
|
||||
A callable that creates PyArrowComputeUDFExpr instances
|
||||
"""
|
||||
|
||||
return _create_pyarrow_compute_udf(pc_func, return_dtype=return_dtype)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _StringNamespace:
|
||||
"""Namespace for string operations on expression columns.
|
||||
|
||||
This namespace provides methods for operating on string-typed columns using
|
||||
PyArrow compute functions.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> # Convert to uppercase
|
||||
>>> expr = col("name").str.upper()
|
||||
>>> # Get string length
|
||||
>>> expr = col("name").str.len()
|
||||
>>> # Check if string starts with a prefix
|
||||
>>> expr = col("name").str.starts_with("A")
|
||||
"""
|
||||
|
||||
_expr: Expr
|
||||
|
||||
# Length methods
|
||||
def len(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Get the length of each string in characters."""
|
||||
return _create_str_udf(pc.utf8_length, DataType.int32())(self._expr)
|
||||
|
||||
def byte_len(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Get the length of each string in bytes."""
|
||||
return _create_str_udf(pc.binary_length, DataType.int32())(self._expr)
|
||||
|
||||
# Case methods
|
||||
def upper(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Convert strings to uppercase."""
|
||||
return _create_str_udf(pc.utf8_upper, DataType.string())(self._expr)
|
||||
|
||||
def lower(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Convert strings to lowercase."""
|
||||
return _create_str_udf(pc.utf8_lower, DataType.string())(self._expr)
|
||||
|
||||
def capitalize(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Capitalize the first character of each string."""
|
||||
return _create_str_udf(pc.utf8_capitalize, DataType.string())(self._expr)
|
||||
|
||||
def title(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Convert strings to title case."""
|
||||
return _create_str_udf(pc.utf8_title, DataType.string())(self._expr)
|
||||
|
||||
def swapcase(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Swap the case of each character."""
|
||||
return _create_str_udf(pc.utf8_swapcase, DataType.string())(self._expr)
|
||||
|
||||
# Predicate methods
|
||||
def is_alpha(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only alphabetic characters."""
|
||||
return _create_str_udf(pc.utf8_is_alpha, DataType.bool())(self._expr)
|
||||
|
||||
def is_alnum(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only alphanumeric characters."""
|
||||
return _create_str_udf(pc.utf8_is_alnum, DataType.bool())(self._expr)
|
||||
|
||||
def is_digit(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only digits."""
|
||||
return _create_str_udf(pc.utf8_is_digit, DataType.bool())(self._expr)
|
||||
|
||||
def is_decimal(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only decimal characters."""
|
||||
return _create_str_udf(pc.utf8_is_decimal, DataType.bool())(self._expr)
|
||||
|
||||
def is_numeric(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only numeric characters."""
|
||||
return _create_str_udf(pc.utf8_is_numeric, DataType.bool())(self._expr)
|
||||
|
||||
def is_space(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only whitespace."""
|
||||
return _create_str_udf(pc.utf8_is_space, DataType.bool())(self._expr)
|
||||
|
||||
def is_lower(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings are lowercase."""
|
||||
return _create_str_udf(pc.utf8_is_lower, DataType.bool())(self._expr)
|
||||
|
||||
def is_upper(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings are uppercase."""
|
||||
return _create_str_udf(pc.utf8_is_upper, DataType.bool())(self._expr)
|
||||
|
||||
def is_title(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings are title-cased."""
|
||||
return _create_str_udf(pc.utf8_is_title, DataType.bool())(self._expr)
|
||||
|
||||
def is_printable(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only printable characters."""
|
||||
return _create_str_udf(pc.utf8_is_printable, DataType.bool())(self._expr)
|
||||
|
||||
def is_ascii(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain only ASCII characters."""
|
||||
return _create_str_udf(pc.string_is_ascii, DataType.bool())(self._expr)
|
||||
|
||||
# Searching methods
|
||||
def starts_with(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings start with a pattern."""
|
||||
return _create_str_udf(pc.starts_with, DataType.bool())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def ends_with(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings end with a pattern."""
|
||||
return _create_str_udf(pc.ends_with, DataType.bool())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def contains(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings contain a substring."""
|
||||
return _create_str_udf(pc.match_substring, DataType.bool())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def match(self, pattern: str, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Match strings against a SQL LIKE pattern."""
|
||||
return _create_str_udf(pc.match_like, DataType.bool())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def find(self, pattern: str, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Find the first occurrence of a substring."""
|
||||
return _create_str_udf(pc.find_substring, DataType.int32())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def count(self, pattern: str, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Count occurrences of a substring."""
|
||||
return _create_str_udf(pc.count_substring, DataType.int32())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def find_regex(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Find the first occurrence matching a regex pattern."""
|
||||
return _create_str_udf(pc.find_substring_regex, DataType.int32())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def count_regex(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Count occurrences matching a regex pattern."""
|
||||
return _create_str_udf(pc.count_substring_regex, DataType.int32())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def match_regex(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Check if strings match a regex pattern."""
|
||||
return _create_str_udf(pc.match_substring_regex, DataType.bool())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
# Transformation methods
|
||||
def reverse(self) -> "PyArrowComputeUDFExpr":
|
||||
"""Reverse each string."""
|
||||
return _create_str_udf(pc.utf8_reverse, DataType.string())(self._expr)
|
||||
|
||||
def slice(self, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Slice strings by codeunit indices."""
|
||||
return _create_str_udf(pc.utf8_slice_codeunits, DataType.string())(
|
||||
self._expr, *args, **kwargs
|
||||
)
|
||||
|
||||
def replace(
|
||||
self, pattern: str, replacement: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Replace occurrences of a substring."""
|
||||
return _create_str_udf(pc.replace_substring, DataType.string())(
|
||||
self._expr, pattern, replacement, *args, **kwargs
|
||||
)
|
||||
|
||||
def replace_regex(
|
||||
self, pattern: str, replacement: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Replace occurrences matching a regex pattern."""
|
||||
return _create_str_udf(pc.replace_substring_regex, DataType.string())(
|
||||
self._expr, pattern, replacement, *args, **kwargs
|
||||
)
|
||||
|
||||
def replace_slice(
|
||||
self, start: int, stop: int, replacement: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Replace a slice with a string."""
|
||||
return _create_str_udf(pc.binary_replace_slice, DataType.string())(
|
||||
self._expr, start, stop, replacement, *args, **kwargs
|
||||
)
|
||||
|
||||
def split(self, pattern: str, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Split strings by a pattern."""
|
||||
return _create_str_udf(pc.split_pattern, DataType(object))(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def split_regex(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Split strings by a regex pattern."""
|
||||
return _create_str_udf(pc.split_pattern_regex, DataType(object))(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def split_whitespace(self, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Split strings on whitespace."""
|
||||
return _create_str_udf(pc.utf8_split_whitespace, DataType(object))(
|
||||
self._expr, *args, **kwargs
|
||||
)
|
||||
|
||||
def extract(
|
||||
self, pattern: str, *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract a substring matching a regex pattern."""
|
||||
return _create_str_udf(pc.extract_regex, DataType.string())(
|
||||
self._expr, pattern, *args, **kwargs
|
||||
)
|
||||
|
||||
def repeat(self, n: int, *args: Any, **kwargs: Any) -> "PyArrowComputeUDFExpr":
|
||||
"""Repeat each string n times."""
|
||||
return _create_str_udf(pc.binary_repeat, DataType.string())(
|
||||
self._expr, n, *args, **kwargs
|
||||
)
|
||||
|
||||
def center(
|
||||
self, width: int, padding: str = " ", *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Center strings in a field of given width."""
|
||||
return _create_str_udf(pc.utf8_center, DataType.string())(
|
||||
self._expr, width, padding, *args, **kwargs
|
||||
)
|
||||
|
||||
def lpad(
|
||||
self, width: int, padding: str = " ", *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Right-align strings by padding with a given character while respecting ``width``.
|
||||
|
||||
If the string is longer than the specified width, it remains intact (no truncation occurs).
|
||||
"""
|
||||
return _create_str_udf(pc.utf8_lpad, DataType.string())(
|
||||
self._expr, width, padding, *args, **kwargs
|
||||
)
|
||||
|
||||
def rpad(
|
||||
self, width: int, padding: str = " ", *args: Any, **kwargs: Any
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Left-align strings by padding with a given character while respecting ``width``.
|
||||
|
||||
If the string is longer than the specified width, it remains intact (no truncation occurs).
|
||||
"""
|
||||
return _create_str_udf(pc.utf8_rpad, DataType.string())(
|
||||
self._expr, width, padding, *args, **kwargs
|
||||
)
|
||||
|
||||
def strip(self, characters: str | None = None) -> "PyArrowComputeUDFExpr":
|
||||
"""Remove leading and trailing whitespace or specified characters.
|
||||
|
||||
Args:
|
||||
characters: Characters to remove. If None, removes whitespace.
|
||||
|
||||
Returns:
|
||||
Expression that strips characters from both ends.
|
||||
"""
|
||||
if characters is None:
|
||||
return _create_str_udf(pc.utf8_trim_whitespace, DataType.string())(
|
||||
self._expr
|
||||
)
|
||||
return _create_str_udf(pc.utf8_trim, DataType.string())(
|
||||
self._expr, characters=characters
|
||||
)
|
||||
|
||||
def lstrip(self, characters: str | None = None) -> "PyArrowComputeUDFExpr":
|
||||
"""Remove leading whitespace or specified characters.
|
||||
|
||||
Args:
|
||||
characters: Characters to remove. If None, removes whitespace.
|
||||
|
||||
Returns:
|
||||
Expression that strips characters from the left.
|
||||
"""
|
||||
if characters is None:
|
||||
return _create_str_udf(pc.utf8_ltrim_whitespace, DataType.string())(
|
||||
self._expr
|
||||
)
|
||||
return _create_str_udf(pc.utf8_ltrim, DataType.string())(
|
||||
self._expr, characters=characters
|
||||
)
|
||||
|
||||
def rstrip(self, characters: str | None = None) -> "PyArrowComputeUDFExpr":
|
||||
"""Remove trailing whitespace or specified characters.
|
||||
|
||||
Args:
|
||||
characters: Characters to remove. If None, removes whitespace.
|
||||
|
||||
Returns:
|
||||
Expression that strips characters from the right.
|
||||
"""
|
||||
if characters is None:
|
||||
return _create_str_udf(pc.utf8_rtrim_whitespace, DataType.string())(
|
||||
self._expr
|
||||
)
|
||||
return _create_str_udf(pc.utf8_rtrim, DataType.string())(
|
||||
self._expr, characters=characters
|
||||
)
|
||||
|
||||
def pad(
|
||||
self,
|
||||
width: int,
|
||||
fillchar: str = " ",
|
||||
side: Literal["left", "right", "both"] = "right",
|
||||
) -> "PyArrowComputeUDFExpr":
|
||||
"""Pad strings to a specified width.
|
||||
|
||||
Args:
|
||||
width: Target width.
|
||||
fillchar: Character to use for padding.
|
||||
side: "left", "right", or "both" for padding side.
|
||||
|
||||
Returns:
|
||||
Expression that pads strings to the given width.
|
||||
"""
|
||||
if side == "right":
|
||||
pc_func = pc.utf8_rpad
|
||||
elif side == "left":
|
||||
pc_func = pc.utf8_lpad
|
||||
elif side == "both":
|
||||
pc_func = pc.utf8_center
|
||||
else:
|
||||
raise ValueError("side must be 'left', 'right', or 'both'")
|
||||
return _create_str_udf(pc_func, DataType.string())(
|
||||
self._expr, width=width, padding=fillchar
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Struct namespace for expression operations on struct-typed columns."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
import pyarrow
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.expressions import _create_pyarrow_compute_udf
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.expressions import Expr, PyArrowComputeUDFExpr
|
||||
|
||||
|
||||
@dataclass
|
||||
class _StructNamespace:
|
||||
"""Namespace for struct operations on expression columns.
|
||||
|
||||
This namespace provides methods for operating on struct-typed columns using
|
||||
PyArrow compute functions.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> # Access a field using method
|
||||
>>> expr = col("user_record").struct.field("age")
|
||||
>>> # Access a field using bracket notation
|
||||
>>> expr = col("user_record").struct["age"]
|
||||
>>> # Access nested field
|
||||
>>> expr = col("user_record").struct["address"].struct["city"]
|
||||
"""
|
||||
|
||||
_expr: Expr
|
||||
|
||||
def __getitem__(self, key: Union[str, int]) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract a field using bracket notation.
|
||||
|
||||
Args:
|
||||
key: The field name or index to extract.
|
||||
|
||||
Returns:
|
||||
PyArrowComputeUDFExpr that extracts the specified field from each struct.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> expr = col("user").struct["age"] # Get age field by name
|
||||
>>> expr = col("user").struct[1] # Get second field by index
|
||||
>>> expr = col("user").struct["address"].struct["city"] # Get nested city field
|
||||
"""
|
||||
if isinstance(key, str):
|
||||
return self.field(key)
|
||||
if isinstance(key, int) and not isinstance(key, bool):
|
||||
return self.field_by_index(key)
|
||||
raise TypeError(
|
||||
f"Struct indices must be strings or integers, not {type(key).__name__}"
|
||||
)
|
||||
|
||||
def field(self, field_name: str) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract a field from a struct.
|
||||
|
||||
Args:
|
||||
field_name: The name of the field to extract.
|
||||
|
||||
Returns:
|
||||
UDFExpr that extracts the specified field from each struct.
|
||||
"""
|
||||
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_struct(arrow_type):
|
||||
try:
|
||||
field_type = arrow_type.field(field_name).type
|
||||
return_dtype = DataType.from_arrow(field_type)
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
return _create_pyarrow_compute_udf(pc.struct_field, return_dtype)(
|
||||
self._expr, field_name
|
||||
)
|
||||
|
||||
def field_by_index(self, index: int) -> "PyArrowComputeUDFExpr":
|
||||
"""Extract a field from a struct by index.
|
||||
|
||||
Args:
|
||||
index: The index of the field to extract.
|
||||
|
||||
Returns:
|
||||
UDFExpr that extracts the specified field from each struct.
|
||||
"""
|
||||
if not isinstance(index, int) or isinstance(index, bool):
|
||||
raise TypeError(
|
||||
f"Struct field index must be an integer, not {type(index).__name__}"
|
||||
)
|
||||
if index < 0:
|
||||
raise ValueError(f"Struct field index must be non-negative, got {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_struct(arrow_type):
|
||||
try:
|
||||
field_type = arrow_type[index].type
|
||||
return_dtype = DataType.from_arrow(field_type)
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
return _create_pyarrow_compute_udf(pc.struct_field, return_dtype)(
|
||||
self._expr, index
|
||||
)
|
||||
Reference in New Issue
Block a user