wehub-resource-sync
695 lines
24 KiB
Python
695 lines
24 KiB
Python
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
|