chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,694 @@
|
||||
import logging
|
||||
import random
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Tuple,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
from ray._common.utils import env_integer
|
||||
from ray.data._internal.arrow_ops import transform_polars, transform_pyarrow
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import shuffle
|
||||
from ray.data._internal.row import row_repr, row_repr_pretty, row_str
|
||||
from ray.data._internal.table_block import TableBlockAccessor, TableBlockBuilder
|
||||
from ray.data._internal.tensor_extensions.arrow import (
|
||||
convert_to_pyarrow_array,
|
||||
pyarrow_table_from_pydict,
|
||||
)
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockAccessor,
|
||||
BlockColumn,
|
||||
BlockColumnAccessor,
|
||||
BlockExecStats,
|
||||
BlockMetadataWithSchema,
|
||||
BlockType,
|
||||
U,
|
||||
)
|
||||
from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE, DataContext
|
||||
from ray.data.expressions import Expr
|
||||
|
||||
try:
|
||||
import pyarrow
|
||||
except ImportError:
|
||||
pyarrow = None
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pandas
|
||||
|
||||
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY = parse_version("13.0.0")
|
||||
_BATCH_SIZE_PRESERVING_STUB_COL_NAME = "__bsp_stub"
|
||||
|
||||
|
||||
def _is_user_visible_column(name: str) -> bool:
|
||||
return name != _BATCH_SIZE_PRESERVING_STUB_COL_NAME
|
||||
|
||||
|
||||
# Set the max chunk size in bytes for Arrow to Batches conversion in
|
||||
# ArrowBlockAccessor.iter_rows(). Default to 4MB, to optimize for image
|
||||
# datasets in parquet format.
|
||||
ARROW_MAX_CHUNK_SIZE_BYTES = env_integer(
|
||||
"RAY_DATA_ARROW_MAX_CHUNK_SIZE_BYTES",
|
||||
int(DEFAULT_TARGET_MAX_BLOCK_SIZE / 32),
|
||||
)
|
||||
|
||||
|
||||
# We offload some transformations to polars for performance.
|
||||
def get_sort_transform(context: DataContext) -> Callable:
|
||||
if context.use_polars or context.use_polars_sort:
|
||||
return transform_polars.sort
|
||||
else:
|
||||
return transform_pyarrow.sort
|
||||
|
||||
|
||||
def get_concat_and_sort_transform(context: DataContext) -> Callable:
|
||||
if context.use_polars or context.use_polars_sort:
|
||||
return transform_polars.concat_and_sort
|
||||
else:
|
||||
return transform_pyarrow.concat_and_sort
|
||||
|
||||
|
||||
class ArrowRow(Mapping):
|
||||
"""
|
||||
Row of a tabular Dataset backed by a Arrow Table block.
|
||||
"""
|
||||
|
||||
def __init__(self, row: Any):
|
||||
self._row = row
|
||||
|
||||
def __getitem__(self, key: Union[str, List[str]]) -> Any:
|
||||
from ray.data.extensions import get_arrow_extension_tensor_types
|
||||
|
||||
tensor_arrow_extension_types = get_arrow_extension_tensor_types()
|
||||
|
||||
def get_item(keys: List[str]) -> Any:
|
||||
schema = self._row.schema
|
||||
if isinstance(schema.field(keys[0]).type, tensor_arrow_extension_types):
|
||||
# Build a tensor row.
|
||||
return tuple(
|
||||
[
|
||||
ArrowBlockAccessor._build_tensor_row(
|
||||
self._row, col_name=key, row_idx=0
|
||||
)
|
||||
for key in keys
|
||||
]
|
||||
)
|
||||
|
||||
table = self._row.select(keys)
|
||||
if len(table) == 0:
|
||||
return None
|
||||
|
||||
items = [col[0] for col in table.columns]
|
||||
try:
|
||||
# Try to interpret this as a pyarrow.Scalar value.
|
||||
return tuple([item.as_py() for item in items])
|
||||
|
||||
except AttributeError:
|
||||
# Assume that this row is an element of an extension array, and
|
||||
# that it is bypassing pyarrow's scalar model for Arrow < 8.0.0.
|
||||
return items
|
||||
|
||||
is_single_item = isinstance(key, str)
|
||||
keys = [key] if is_single_item else key
|
||||
|
||||
items = get_item(keys)
|
||||
|
||||
if items is None:
|
||||
return None
|
||||
elif is_single_item:
|
||||
return items[0]
|
||||
else:
|
||||
return items
|
||||
|
||||
def __iter__(self) -> Iterator:
|
||||
for k in self._row.column_names:
|
||||
yield k
|
||||
|
||||
def __len__(self):
|
||||
return self._row.num_columns
|
||||
|
||||
def as_pydict(self) -> Dict[str, Any]:
|
||||
return dict(self.items())
|
||||
|
||||
def __str__(self):
|
||||
return row_str(self)
|
||||
|
||||
def __repr__(self):
|
||||
return row_repr(self)
|
||||
|
||||
def _repr_pretty_(self, p, cycle):
|
||||
return row_repr_pretty(self, p, cycle)
|
||||
|
||||
|
||||
class ArrowBlockBuilder(TableBlockBuilder):
|
||||
def __init__(self):
|
||||
if pyarrow is None:
|
||||
raise ImportError("Run `pip install pyarrow` for Arrow support")
|
||||
super().__init__((pyarrow.Table, bytes))
|
||||
|
||||
@staticmethod
|
||||
def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:
|
||||
return pyarrow_table_from_pydict(
|
||||
{
|
||||
column_name: convert_to_pyarrow_array(column_values, column_name)
|
||||
for column_name, column_values in columns.items()
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _combine_tables(tables: List[Block]) -> Block:
|
||||
if len(tables) > 1:
|
||||
return transform_pyarrow.concat(
|
||||
tables, promote_types=True, preserve_order=True
|
||||
)
|
||||
else:
|
||||
return tables[0]
|
||||
|
||||
@staticmethod
|
||||
def _concat_would_copy() -> bool:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _empty_table() -> "pyarrow.Table":
|
||||
return pyarrow_table_from_pydict({})
|
||||
|
||||
def block_type(self) -> BlockType:
|
||||
return BlockType.ARROW
|
||||
|
||||
|
||||
def _get_max_chunk_size(
|
||||
table: "pyarrow.Table", max_chunk_size_bytes: int
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
Calculate the max chunk size in rows for Arrow to Batches conversion in
|
||||
ArrowBlockAccessor.iter_rows().
|
||||
Args:
|
||||
table: The pyarrow table to calculate the max chunk size for.
|
||||
max_chunk_size_bytes: The max chunk size in bytes.
|
||||
Returns:
|
||||
The max chunk size in rows, or None if the table is empty.
|
||||
"""
|
||||
if table.nbytes == 0:
|
||||
return None
|
||||
else:
|
||||
avg_row_size = table.nbytes / table.num_rows
|
||||
return max(1, int(max_chunk_size_bytes / avg_row_size))
|
||||
|
||||
|
||||
class ArrowBlockAccessor(TableBlockAccessor):
|
||||
ROW_TYPE = ArrowRow
|
||||
|
||||
def __init__(self, table: "pyarrow.Table"):
|
||||
if pyarrow is None:
|
||||
raise ImportError("Run `pip install pyarrow` for Arrow support")
|
||||
super().__init__(table)
|
||||
self._max_chunk_size: Optional[int] = None
|
||||
|
||||
def _get_row(self, index: int) -> ArrowRow:
|
||||
base_row = self.slice(index, index + 1, copy=False)
|
||||
return ArrowRow(base_row)
|
||||
|
||||
def column_names(self) -> List[str]:
|
||||
return self._table.column_names
|
||||
|
||||
def fill_column(self, name: str, value: Any) -> Block:
|
||||
import pyarrow.compute as pc
|
||||
|
||||
# Check if value is array-like - if so, use upsert_column logic
|
||||
if isinstance(value, (pyarrow.Array, pyarrow.ChunkedArray)):
|
||||
return self.upsert_column(name, value)
|
||||
else:
|
||||
# Scalar value - use original fill_column logic
|
||||
if isinstance(value, pyarrow.Scalar):
|
||||
type = value.type
|
||||
else:
|
||||
type = pyarrow.infer_type([value])
|
||||
|
||||
array = pyarrow.nulls(len(self._table), type=type)
|
||||
array = pc.fill_null(array, value)
|
||||
return self.upsert_column(name, array)
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes) -> "ArrowBlockAccessor":
|
||||
reader = pyarrow.ipc.open_stream(data)
|
||||
return cls(reader.read_all())
|
||||
|
||||
@staticmethod
|
||||
def _build_tensor_row(row: ArrowRow, row_idx: int, col_name: str) -> np.ndarray:
|
||||
|
||||
element = row[col_name][row_idx]
|
||||
arr = element.as_py()
|
||||
|
||||
assert isinstance(arr, np.ndarray), type(arr)
|
||||
return arr
|
||||
|
||||
def slice(self, start: int, end: int, copy: bool = False) -> "pyarrow.Table":
|
||||
view = self._table.slice(start, end - start)
|
||||
if copy:
|
||||
view = transform_pyarrow.combine_chunks(view, copy)
|
||||
return view
|
||||
|
||||
def random_shuffle(self, random_seed: Optional[int]) -> "pyarrow.Table":
|
||||
return shuffle(self._table, random_seed)
|
||||
|
||||
def schema(self) -> "pyarrow.lib.Schema":
|
||||
return self._table.schema
|
||||
|
||||
def to_pandas(self) -> "pandas.DataFrame":
|
||||
from ray.data.util.data_batch_conversion import _cast_tensor_columns_to_ndarrays
|
||||
|
||||
# We specify ignore_metadata=True because pyarrow will use the metadata
|
||||
# to build the Table. This is handled incorrectly for older pyarrow versions
|
||||
ctx = DataContext.get_current()
|
||||
|
||||
# types_mapper preserves Arrow dtypes through the pandas round-trip:
|
||||
# - Standard Arrow types become pd.ArrowDtype, so pa.Table.from_pandas()
|
||||
# can reconstruct them exactly without lossy numpy conversion.
|
||||
# - Extension types (Ray's ArrowTensorType / ArrowPythonObjectType and
|
||||
# pyarrow's native FixedShapeTensorType) return None, falling back to
|
||||
# their own to_pandas_dtype() hooks. Note: native FixedShapeTensorType
|
||||
# subclasses BaseExtensionType but not ExtensionType, so we check the
|
||||
# broader BaseExtensionType.
|
||||
def _types_mapper(t):
|
||||
if isinstance(t, pyarrow.BaseExtensionType) or pyarrow.types.is_dictionary(
|
||||
t
|
||||
):
|
||||
return None
|
||||
return pd.ArrowDtype(t)
|
||||
|
||||
df = self._table.to_pandas(
|
||||
ignore_metadata=ctx.pandas_block_ignore_metadata,
|
||||
types_mapper=_types_mapper,
|
||||
)
|
||||
if ctx.enable_tensor_extension_casting:
|
||||
df = _cast_tensor_columns_to_ndarrays(df, arrow_schema=self._table.schema)
|
||||
return df
|
||||
|
||||
def to_numpy(
|
||||
self, columns: Optional[Union[str, List[str]]] = None
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
if columns is None:
|
||||
columns = self._table.column_names
|
||||
should_be_single_ndarray = False
|
||||
elif isinstance(columns, list):
|
||||
should_be_single_ndarray = False
|
||||
else:
|
||||
columns = [columns]
|
||||
should_be_single_ndarray = True
|
||||
|
||||
column_names_set = set(self._table.column_names)
|
||||
for column in columns:
|
||||
if column not in column_names_set:
|
||||
raise ValueError(
|
||||
f"Cannot find column {column}, available columns: "
|
||||
f"{column_names_set}"
|
||||
)
|
||||
|
||||
column_values_ndarrays = []
|
||||
|
||||
for col_name in columns:
|
||||
col = self._table[col_name]
|
||||
|
||||
# Combine columnar values arrays to make these contiguous
|
||||
# (making them compatible with numpy format)
|
||||
combined_array = transform_pyarrow.combine_chunked_array(col)
|
||||
|
||||
column_values_ndarrays.append(
|
||||
transform_pyarrow.to_numpy(combined_array, zero_copy_only=False)
|
||||
)
|
||||
|
||||
if should_be_single_ndarray:
|
||||
assert len(columns) == 1
|
||||
return column_values_ndarrays[0]
|
||||
else:
|
||||
return dict(zip(columns, column_values_ndarrays))
|
||||
|
||||
def to_arrow(self) -> "pyarrow.Table":
|
||||
return self._table
|
||||
|
||||
def num_rows(self) -> int:
|
||||
# Arrow may represent an empty table via an N > 0 row, 0-column table, e.g. when
|
||||
# slicing an empty table, so we return 0 if num_columns == 0.
|
||||
return self._table.num_rows if self._table.num_columns > 0 else 0
|
||||
|
||||
def size_bytes(self) -> int:
|
||||
return self._table.nbytes
|
||||
|
||||
def _zip(self, acc: BlockAccessor) -> "Block":
|
||||
r = self.to_arrow()
|
||||
s = acc.to_arrow()
|
||||
for col_name in s.column_names:
|
||||
col = s.column(col_name)
|
||||
# Ensure the column names are unique after zip.
|
||||
if col_name in r.column_names:
|
||||
i = 1
|
||||
new_name = col_name
|
||||
while new_name in r.column_names:
|
||||
new_name = "{}_{}".format(col_name, i)
|
||||
i += 1
|
||||
col_name = new_name
|
||||
r = r.append_column(col_name, col)
|
||||
return r
|
||||
|
||||
def upsert_column(
|
||||
self, column_name: str, column_data: BlockColumn
|
||||
) -> "pyarrow.Table":
|
||||
assert isinstance(
|
||||
column_data, (pyarrow.Array, pyarrow.ChunkedArray)
|
||||
), f"Expected either a pyarrow.Array or pyarrow.ChunkedArray, got: {type(column_data)}"
|
||||
|
||||
column_idx = self._table.schema.get_field_index(column_name)
|
||||
if column_idx == -1:
|
||||
return self._table.append_column(column_name, column_data)
|
||||
else:
|
||||
return self._table.set_column(column_idx, column_name, column_data)
|
||||
|
||||
@staticmethod
|
||||
def builder() -> ArrowBlockBuilder:
|
||||
return ArrowBlockBuilder()
|
||||
|
||||
@staticmethod
|
||||
def _empty_table() -> "pyarrow.Table":
|
||||
return ArrowBlockBuilder._empty_table()
|
||||
|
||||
def take(
|
||||
self,
|
||||
indices: Union[List[int], "pyarrow.Array", "pyarrow.ChunkedArray"],
|
||||
) -> "pyarrow.Table":
|
||||
"""Select rows from the underlying table.
|
||||
|
||||
This method is an alternative to pyarrow.Table.take(), which breaks for
|
||||
extension arrays.
|
||||
"""
|
||||
return transform_pyarrow.take_table(self._table, indices)
|
||||
|
||||
def drop(self, columns: List[str]) -> Block:
|
||||
return self._table.drop(columns)
|
||||
|
||||
def select(self, columns: List[str]) -> "pyarrow.Table":
|
||||
if not all(isinstance(col, str) for col in columns):
|
||||
raise ValueError(
|
||||
"Columns must be a list of column name strings when aggregating on "
|
||||
f"Arrow blocks, but got: {columns}."
|
||||
)
|
||||
if len(columns) == 0:
|
||||
# Empty projection (e.g. count or ``select_columns([])``).
|
||||
# Drop every existing column, then append the stub so row
|
||||
# counts survive downstream ``pa.concat_tables`` calls (which
|
||||
# collapse num_rows to 0 when all inputs have 0 columns).
|
||||
# ``pa.Table`` tracks num_rows as metadata independent of
|
||||
# columns, so ``select([])`` preserves it here. The stub is
|
||||
# filtered out of the user-visible schema; it's a physical
|
||||
# placeholder only.
|
||||
narrowed = self._table.select([])
|
||||
return ArrowBlockAccessor(narrowed).fill_column(
|
||||
_BATCH_SIZE_PRESERVING_STUB_COL_NAME, None
|
||||
)
|
||||
return self._table.select(columns)
|
||||
|
||||
def rename_columns(self, columns_rename: Dict[str, str]) -> "pyarrow.Table":
|
||||
return self._table.rename_columns(columns_rename)
|
||||
|
||||
def hstack(self, other_block: "pyarrow.Table") -> "pyarrow.Table":
|
||||
|
||||
result_table = self._table
|
||||
for name, column in zip(other_block.column_names, other_block.columns):
|
||||
result_table = result_table.append_column(name, column)
|
||||
|
||||
return result_table
|
||||
|
||||
def _sample(self, n_samples: int, sort_key: "SortKey") -> "pyarrow.Table":
|
||||
indices = random.sample(range(self._table.num_rows), n_samples)
|
||||
table = self._table.select(sort_key.get_columns())
|
||||
return transform_pyarrow.take_table(table, indices)
|
||||
|
||||
def sort(self, sort_key: "SortKey") -> Block:
|
||||
assert (
|
||||
sort_key.get_columns()
|
||||
), f"Sorting columns couldn't be empty (got {sort_key.get_columns()})"
|
||||
|
||||
if self._table.num_rows == 0:
|
||||
# If the pyarrow table is empty we may not have schema
|
||||
# so calling sort_indices() will raise an error.
|
||||
return self._empty_table()
|
||||
|
||||
context = DataContext.get_current()
|
||||
sort = get_sort_transform(context)
|
||||
|
||||
return sort(self._table, sort_key)
|
||||
|
||||
def sort_and_partition(
|
||||
self, boundaries: List[T], sort_key: "SortKey"
|
||||
) -> List["Block"]:
|
||||
table = self.sort(sort_key)
|
||||
|
||||
if table.num_rows == 0:
|
||||
return [self._empty_table() for _ in range(len(boundaries) + 1)]
|
||||
elif len(boundaries) == 0:
|
||||
return [table]
|
||||
|
||||
return BlockAccessor.for_block(table)._find_partitions_sorted(
|
||||
boundaries, sort_key
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def merge_sorted_blocks(
|
||||
blocks: List[Block], sort_key: "SortKey"
|
||||
) -> Tuple[Block, BlockMetadataWithSchema]:
|
||||
stats = BlockExecStats.builder()
|
||||
blocks = [b for b in blocks if b.num_rows > 0]
|
||||
if len(blocks) == 0:
|
||||
ret = ArrowBlockAccessor._empty_table()
|
||||
else:
|
||||
# Handle blocks of different types.
|
||||
blocks = TableBlockAccessor.normalize_block_types(blocks, BlockType.ARROW)
|
||||
concat_and_sort = get_concat_and_sort_transform(DataContext.get_current())
|
||||
ret = concat_and_sort(blocks, sort_key, promote_types=True)
|
||||
return ret, BlockMetadataWithSchema.from_block(
|
||||
ret, block_exec_stats=stats.build()
|
||||
)
|
||||
|
||||
def block_type(self) -> BlockType:
|
||||
return BlockType.ARROW
|
||||
|
||||
def iter_rows(
|
||||
self, public_row_format: bool
|
||||
) -> Iterator[Union[Mapping, np.ndarray]]:
|
||||
table = self._table
|
||||
if public_row_format:
|
||||
from ray.data._internal.utils.transform_pyarrow import (
|
||||
_is_native_tensor_type,
|
||||
)
|
||||
|
||||
if self._max_chunk_size is None:
|
||||
# Calling _get_max_chunk_size in constructor makes it slow, so we
|
||||
# are calling it here only when needed.
|
||||
self._max_chunk_size = _get_max_chunk_size(
|
||||
table, ARROW_MAX_CHUNK_SIZE_BYTES
|
||||
)
|
||||
contains_native_tensor_columns = any(
|
||||
_is_native_tensor_type(column.type) for column in table.columns
|
||||
)
|
||||
for batch in table.to_batches(max_chunksize=self._max_chunk_size):
|
||||
if contains_native_tensor_columns:
|
||||
# HACK: For v1 and v2 tensors, we can control what is returned
|
||||
# by overriding ExtensionScalar.as_py (see ArrowTensorScalar).
|
||||
# For pyarrow native FixedShapeTensorArrays we cannot, so we
|
||||
# use _iter_rows_from_batch_with_tensors to handle conversion.
|
||||
yield from _iter_rows_from_batch_with_tensors(batch)
|
||||
else:
|
||||
yield from batch.to_pylist()
|
||||
else:
|
||||
num_rows = self.num_rows()
|
||||
for i in range(num_rows):
|
||||
yield self._get_row(i)
|
||||
|
||||
def filter(self, predicate_expr: "Expr") -> "pyarrow.Table":
|
||||
"""Filter rows based on a predicate expression."""
|
||||
if self._table.num_rows == 0:
|
||||
return self._table
|
||||
|
||||
from ray.data._internal.planner.plan_expression.expression_evaluator import (
|
||||
eval_expr,
|
||||
)
|
||||
|
||||
# Evaluate the expression to get a boolean mask
|
||||
mask = eval_expr(predicate_expr, self._table)
|
||||
|
||||
# Use PyArrow's built-in filter method
|
||||
return self._table.filter(mask)
|
||||
|
||||
|
||||
def _iter_rows_from_batch_with_tensors(
|
||||
batch: "pyarrow.RecordBatch",
|
||||
) -> Iterator[Dict[str, Any]]:
|
||||
"""Iterate over rows in a batch that may contain native tensor columns.
|
||||
|
||||
For pyarrow native FixedShapeTensorArrays, we must manually convert them
|
||||
to ndarrays which preserve shape/ndim. Without this, FixedShapeTensorArrays
|
||||
would be translated to contiguous 1d arrays.
|
||||
|
||||
See: https://arrow.apache.org/docs/python/generated/pyarrow.FixedShapeTensorArray.html
|
||||
|
||||
Args:
|
||||
batch: A PyArrow RecordBatch that may contain tensor columns.
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: Dictionaries mapping column names to values for each row.
|
||||
"""
|
||||
from ray.data._internal.utils.transform_pyarrow import _is_native_tensor_type
|
||||
|
||||
col_values = []
|
||||
for column in batch.columns:
|
||||
if _is_native_tensor_type(column.type):
|
||||
col_values.append(column.to_numpy_ndarray())
|
||||
else:
|
||||
col_values.append(column.to_pylist())
|
||||
|
||||
for idx in range(batch.num_rows):
|
||||
yield {name: col[idx] for name, col in zip(batch.column_names, col_values)}
|
||||
|
||||
|
||||
class ArrowBlockColumnAccessor(BlockColumnAccessor):
|
||||
def __init__(self, col: Union["pyarrow.Array", "pyarrow.ChunkedArray"]):
|
||||
super().__init__(col)
|
||||
|
||||
def count(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
res = pac.count(self._column, mode="only_valid" if ignore_nulls else "all")
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def sum(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
res = pac.sum(self._column, skip_nulls=ignore_nulls)
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def min(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
res = pac.min(self._column, skip_nulls=ignore_nulls)
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def max(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
res = pac.max(self._column, skip_nulls=ignore_nulls)
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def mean(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
res = pac.mean(self._column, skip_nulls=ignore_nulls)
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def sum_of_squared_diffs_from_mean(
|
||||
self, ignore_nulls: bool, mean: Optional[U] = None, as_py: bool = True
|
||||
) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
# Calculate mean if not provided
|
||||
if mean is None:
|
||||
mean = self.mean(ignore_nulls=ignore_nulls)
|
||||
|
||||
if mean is None:
|
||||
return None
|
||||
|
||||
res = pac.sum(
|
||||
pac.power(pac.subtract(self._column, mean), 2), skip_nulls=ignore_nulls
|
||||
)
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def quantile(
|
||||
self, *, q: float, ignore_nulls: bool, as_py: bool = True
|
||||
) -> Optional[U]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
array = pac.quantile(self._column, q=q, skip_nulls=ignore_nulls)
|
||||
# NOTE: That quantile method still returns an array
|
||||
res = array[0]
|
||||
return res.as_py() if as_py else res
|
||||
|
||||
def unique(self) -> BlockColumn:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
if self.is_composed_of_lists():
|
||||
# NOTE: Arrow doesn't provide unique kernels for `ListArray`s and
|
||||
# such, so we rely on Polars to encode and compute unique
|
||||
# values instead
|
||||
import polars
|
||||
|
||||
return polars.from_arrow(self._column).unique().to_arrow()
|
||||
|
||||
return pac.unique(self._column)
|
||||
|
||||
def value_counts(self) -> Optional[Dict[str, List]]:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
value_counts: pyarrow.StructArray = pac.value_counts(self._column)
|
||||
if len(value_counts) == 0:
|
||||
return None
|
||||
return {
|
||||
"values": value_counts.field("values").to_pylist(),
|
||||
"counts": value_counts.field("counts").to_pylist(),
|
||||
}
|
||||
|
||||
def hash(self) -> BlockColumn:
|
||||
import polars as pl
|
||||
|
||||
df = pl.DataFrame({"col": self._column})
|
||||
hashes = df.hash_rows().cast(pl.Int64, wrap_numerical=True)
|
||||
return hashes.to_arrow()
|
||||
|
||||
def flatten(self) -> BlockColumn:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
return pac.list_flatten(self._column)
|
||||
|
||||
def dropna(self) -> BlockColumn:
|
||||
import pyarrow.compute as pac
|
||||
|
||||
return pac.drop_null(self._column)
|
||||
|
||||
def is_composed_of_lists(self) -> bool:
|
||||
types = (pyarrow.lib.ListType, pyarrow.lib.LargeListType)
|
||||
return isinstance(self._column.type, types)
|
||||
|
||||
def to_pylist(self) -> List[Any]:
|
||||
return self._column.to_pylist()
|
||||
|
||||
def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray:
|
||||
if get_pyarrow_version() < _MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY:
|
||||
if isinstance(
|
||||
self._column, pyarrow.ChunkedArray
|
||||
): # NOTE: ChunkedArray in Pyarrow < 13.0.0 does not support ``zero_copy_only``
|
||||
return self._column.to_numpy()
|
||||
else:
|
||||
return self._column.to_numpy(zero_copy_only=zero_copy_only)
|
||||
|
||||
return self._column.to_numpy(zero_copy_only=zero_copy_only)
|
||||
|
||||
def _to_arrow_compatible_container(self) -> Union[List[Any], "pyarrow.Array"]:
|
||||
return self._column
|
||||
Reference in New Issue
Block a user