777 lines
28 KiB
Python
777 lines
28 KiB
Python
import collections
|
|
import logging
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Mapping,
|
|
Optional,
|
|
Tuple,
|
|
TypeVar,
|
|
Union,
|
|
)
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pandas.api.types import is_object_dtype, is_scalar, is_string_dtype
|
|
|
|
from ray.data._internal.numpy_support import convert_to_numpy
|
|
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.utils import _should_convert_to_tensor
|
|
from ray.data._internal.util import is_null
|
|
from ray.data.block import (
|
|
Block,
|
|
BlockAccessor,
|
|
BlockColumn,
|
|
BlockColumnAccessor,
|
|
BlockExecStats,
|
|
BlockType,
|
|
U,
|
|
)
|
|
from ray.data.context import DataContext
|
|
from ray.data.expressions import Expr
|
|
|
|
if TYPE_CHECKING:
|
|
import pandas
|
|
import pyarrow
|
|
|
|
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
|
from ray.data.block import BlockMetadataWithSchema
|
|
|
|
T = TypeVar("T")
|
|
# Max number of samples used to estimate the Pandas block size.
|
|
_PANDAS_SIZE_BYTES_MAX_SAMPLE_COUNT = 200
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_pandas = None
|
|
|
|
|
|
def lazy_import_pandas():
|
|
global _pandas
|
|
if _pandas is None:
|
|
import pandas
|
|
|
|
_pandas = pandas
|
|
return _pandas
|
|
|
|
|
|
def _from_pandas_safe(df: "pandas.DataFrame") -> "pyarrow.Table":
|
|
"""Convert a pandas DataFrame to an Arrow table, handling object-dtype columns.
|
|
|
|
``pa.Table.from_pandas`` infers Arrow types for object-dtype columns by inspecting
|
|
the first Python value, then calls ``pa.array()`` on the whole column. This fails
|
|
for values that PyArrow cannot convert natively — e.g. multi-dimensional numpy
|
|
arrays, PIL images, or mixed list/scalar content.
|
|
|
|
This function routes object-dtype columns through ``convert_to_pyarrow_array``,
|
|
which produces ``ArrowTensorArray`` for ndarray elements and falls back to
|
|
``ArrowPythonObjectArray`` (pickle) for arbitrary Python objects. All other columns
|
|
go through ``pa.array(col, from_pandas=True)`` which handles nullable dtypes and
|
|
extension types via ``__arrow_array__``.
|
|
"""
|
|
import pyarrow as pa
|
|
|
|
from ray.data._internal.tensor_extensions.arrow import convert_to_pyarrow_array
|
|
|
|
# If no object-dtype columns, use fast path with regular from_pandas()
|
|
if not any(is_object_dtype(df[col].dtype) for col in df.columns):
|
|
# Set `preserve_index=False` so that Arrow doesn't add a '__index_level_0__'
|
|
return pa.Table.from_pandas(df, preserve_index=False)
|
|
|
|
# Convert column by column: object-dtype columns go through
|
|
# convert_to_pyarrow_array (handles tensors, PIL images, arbitrary objects),
|
|
# all others go through pa.array() with from_pandas=True.
|
|
arrays = []
|
|
fields = []
|
|
for col_name in df.columns:
|
|
col = df[col_name]
|
|
if is_object_dtype(col.dtype):
|
|
arr = convert_to_pyarrow_array(col.values, col_name)
|
|
else:
|
|
arr = pa.array(col, from_pandas=True)
|
|
arrays.append(arr)
|
|
fields.append(pa.field(col_name, arr.type))
|
|
|
|
return pa.table(dict(zip(df.columns, arrays)), schema=pa.schema(fields))
|
|
|
|
|
|
class PandasRow(Mapping):
|
|
"""
|
|
Row of a tabular Dataset backed by a Pandas DataFrame block.
|
|
"""
|
|
|
|
def __init__(self, row: Any):
|
|
self._row = row
|
|
|
|
def __getitem__(self, key: Union[str, List[str]]) -> Any:
|
|
from ray.data.extensions import TensorArrayElement
|
|
|
|
def get_item(keys: List[str]) -> Any:
|
|
col = self._row[keys]
|
|
if len(col) == 0:
|
|
return None
|
|
|
|
items = col.iloc[0]
|
|
if isinstance(items.iloc[0], TensorArrayElement):
|
|
# Getting an item in a Pandas tensor column may return
|
|
# a TensorArrayElement, which we have to convert to an ndarray.
|
|
return tuple(item.to_numpy() for item in items)
|
|
|
|
try:
|
|
# Try to interpret this as a numpy-type value.
|
|
# See https://stackoverflow.com/questions/9452775/converting-numpy-dtypes-to-native-python-types. # noqa: E501
|
|
return tuple(item for item in items)
|
|
|
|
except (AttributeError, ValueError) as e:
|
|
logger.warning(f"Failed to convert {items} to a tuple", exc_info=e)
|
|
|
|
# Fallback to the original form.
|
|
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.columns:
|
|
yield k
|
|
|
|
def __len__(self):
|
|
return self._row.shape[1]
|
|
|
|
def as_pydict(self) -> Dict[str, Any]:
|
|
pydict: Dict[str, Any] = {}
|
|
for key, value in self.items():
|
|
# Convert NA to None for consistency across block formats. `pd.isna`
|
|
# returns True for both NA and NaN, but since we want to preserve NaN
|
|
# values, we check for identity instead.
|
|
if is_scalar(value) and value is pd.NA:
|
|
pydict[key] = None
|
|
else:
|
|
pydict[key] = value
|
|
|
|
return pydict
|
|
|
|
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 PandasBlockColumnAccessor(BlockColumnAccessor):
|
|
def __init__(self, col: "pandas.Series"):
|
|
super().__init__(col)
|
|
|
|
def count(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
|
return self._column.count() if ignore_nulls else len(self._column)
|
|
|
|
def sum(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
|
# NOTE: Pandas ``Series`` isn't able to properly handle the case with
|
|
# all-null/NaN values in the column, hence we have to handle it here
|
|
if self._is_all_null():
|
|
return None
|
|
|
|
# NOTE: We pass `min_count=1` to workaround quirky Pandas behavior,
|
|
# where (by default) when min_count=0 it will return 0.0 for
|
|
# all-null/NaN series
|
|
return self._column.sum(skipna=ignore_nulls, min_count=1)
|
|
|
|
def min(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
|
# NOTE: Pandas ``Series`` isn't able to properly handle the case with
|
|
# all-null/NaN values in the column, hence we have to handle it here
|
|
if self._is_all_null():
|
|
return None
|
|
|
|
return self._column.min(skipna=ignore_nulls)
|
|
|
|
def max(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
|
# NOTE: Pandas ``Series`` isn't able to properly handle the case with
|
|
# all-null/NaN values in the column, hence we have to handle it here
|
|
if self._is_all_null():
|
|
return None
|
|
|
|
return self._column.max(skipna=ignore_nulls)
|
|
|
|
def mean(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
|
|
# NOTE: We manually implement mean here to keep implementation consistent
|
|
# with behavior of ``sum`` method returning null if the series
|
|
# contains exclusively null values
|
|
sum_ = self.sum(ignore_nulls=ignore_nulls)
|
|
|
|
return (
|
|
sum_ / self.count(ignore_nulls=ignore_nulls) if not is_null(sum_) else sum_
|
|
)
|
|
|
|
def quantile(
|
|
self, *, q: float, ignore_nulls: bool, as_py: bool = True
|
|
) -> Optional[U]:
|
|
return self._column.quantile(q=q)
|
|
|
|
def value_counts(self) -> Optional[Dict[str, List]]:
|
|
value_counts = self._column.value_counts()
|
|
if len(value_counts) == 0:
|
|
return None
|
|
return {
|
|
"values": value_counts.index.tolist(),
|
|
"counts": value_counts.values.tolist(),
|
|
}
|
|
|
|
def hash(self) -> BlockColumn:
|
|
|
|
from ray.data._internal.tensor_extensions.pandas import TensorArrayElement
|
|
|
|
first_non_null = next((x for x in self._column if x is not None), None)
|
|
if isinstance(first_non_null, TensorArrayElement):
|
|
self._column = self._column.apply(lambda x: x.to_numpy())
|
|
|
|
import polars as pl
|
|
|
|
df = pl.from_pandas(self._column.to_frame())
|
|
hashes = df.hash_rows().cast(pl.Int64, wrap_numerical=True)
|
|
return hashes.to_pandas()
|
|
|
|
def unique(self) -> BlockColumn:
|
|
|
|
pd = lazy_import_pandas()
|
|
|
|
try:
|
|
if self.is_composed_of_lists():
|
|
# NOTE: Pandas uses hashing internally to compute unique values,
|
|
# and hence we have to convert lists into tuples to make
|
|
# them hashable
|
|
col = self._column.map(lambda l: l if l is None else tuple(l))
|
|
else:
|
|
col = self._column
|
|
|
|
return pd.Series(col.unique())
|
|
except ValueError as e:
|
|
if "buffer source array is read-only" in str(e):
|
|
# NOTE: Pandas < 2.0 somehow tries to update the underlying buffer
|
|
# when computing unique values hence failing
|
|
return pd.Series(self._column.copy().unique())
|
|
else:
|
|
raise
|
|
|
|
def flatten(self) -> BlockColumn:
|
|
from ray.data._internal.tensor_extensions.pandas import TensorArrayElement
|
|
|
|
first_non_null = next((x for x in self._column if x is not None), None)
|
|
if not isinstance(first_non_null, TensorArrayElement):
|
|
column = self._column
|
|
else:
|
|
column = self._column.apply(
|
|
lambda x: x.to_numpy() if isinstance(x, TensorArrayElement) else x
|
|
)
|
|
|
|
# NOTE: `Series.explode` explodes empty lists into NaNs, necessitating
|
|
# filtering out of empty lists first
|
|
if self.is_composed_of_lists():
|
|
mask = column.apply(lambda x: x is not None and len(x) > 0)
|
|
column = column[mask]
|
|
|
|
return column.explode(ignore_index=True)
|
|
|
|
def dropna(self) -> BlockColumn:
|
|
return self._column.dropna()
|
|
|
|
def sum_of_squared_diffs_from_mean(
|
|
self,
|
|
ignore_nulls: bool,
|
|
mean: Optional[U] = None,
|
|
as_py: bool = True,
|
|
) -> Optional[U]:
|
|
if mean is None:
|
|
mean = self.mean(ignore_nulls=ignore_nulls)
|
|
|
|
if is_null(mean):
|
|
return mean
|
|
|
|
return ((self._column - mean) ** 2).sum(skipna=ignore_nulls)
|
|
|
|
def to_pylist(self) -> List[Any]:
|
|
return self._column.to_list()
|
|
|
|
def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray:
|
|
"""NOTE: Unlike Arrow, specifying `zero_copy_only=True` isn't a guarantee
|
|
that no copy will be made
|
|
"""
|
|
|
|
return self._column.to_numpy(copy=not zero_copy_only)
|
|
|
|
def _to_arrow_compatible_container(self) -> Union[List[Any], "pyarrow.Array"]:
|
|
return self.to_pylist()
|
|
|
|
def _is_all_null(self):
|
|
return not self._column.notna().any()
|
|
|
|
def is_composed_of_lists(self) -> bool:
|
|
from ray.data._internal.tensor_extensions.pandas import TensorArrayElement
|
|
|
|
types = (list, np.ndarray, TensorArrayElement)
|
|
first_non_null = next((x for x in self._column if x is not None), None)
|
|
return isinstance(first_non_null, types)
|
|
|
|
|
|
class PandasBlockBuilder(TableBlockBuilder):
|
|
def __init__(self):
|
|
pandas = lazy_import_pandas()
|
|
super().__init__(pandas.DataFrame)
|
|
|
|
@staticmethod
|
|
def _table_from_pydict(columns: Dict[str, List[Any]]) -> "pandas.DataFrame":
|
|
from ray.data.extensions.tensor_extension import TensorArray
|
|
|
|
pandas = lazy_import_pandas()
|
|
return pandas.DataFrame(
|
|
{
|
|
column_name: (
|
|
TensorArray(convert_to_numpy(column_values))
|
|
if len(column_values) > 0
|
|
and _should_convert_to_tensor(column_values, column_name)
|
|
else column_values
|
|
)
|
|
for column_name, column_values in columns.items()
|
|
}
|
|
)
|
|
|
|
@staticmethod
|
|
def _combine_tables(tables: List["pandas.DataFrame"]) -> "pandas.DataFrame":
|
|
pandas = lazy_import_pandas()
|
|
from ray.data.util.data_batch_conversion import (
|
|
_cast_ndarray_columns_to_tensor_extension,
|
|
)
|
|
|
|
if len(tables) > 1:
|
|
df = pandas.concat(tables, ignore_index=True)
|
|
df.reset_index(drop=True, inplace=True)
|
|
else:
|
|
df = tables[0]
|
|
|
|
ctx = DataContext.get_current()
|
|
if ctx.enable_tensor_extension_casting:
|
|
df = _cast_ndarray_columns_to_tensor_extension(df)
|
|
|
|
return df
|
|
|
|
@staticmethod
|
|
def _concat_would_copy() -> bool:
|
|
return True
|
|
|
|
@staticmethod
|
|
def _empty_table() -> "pandas.DataFrame":
|
|
pandas = lazy_import_pandas()
|
|
return pandas.DataFrame()
|
|
|
|
def block_type(self) -> BlockType:
|
|
return BlockType.PANDAS
|
|
|
|
|
|
# NOTE: This has to be compatible with Pyarrow ``Schema``
|
|
@dataclass(frozen=True, init=False)
|
|
class PandasBlockSchema:
|
|
# Stored as tuples for hash-ability.
|
|
names: Tuple[str, ...]
|
|
types: Tuple
|
|
|
|
def __init__(self, names, types):
|
|
object.__setattr__(self, "names", tuple(names))
|
|
object.__setattr__(self, "types", tuple(types))
|
|
|
|
|
|
class PandasBlockAccessor(TableBlockAccessor):
|
|
ROW_TYPE = PandasRow
|
|
|
|
def __init__(self, table: "pandas.DataFrame"):
|
|
super().__init__(table)
|
|
|
|
def _get_row(self, index: int) -> PandasRow:
|
|
base_row = self.slice(index, index + 1, copy=False)
|
|
return PandasRow(base_row)
|
|
|
|
def column_names(self) -> List[str]:
|
|
return self._table.columns.tolist()
|
|
|
|
def fill_column(self, name: str, value: Any) -> Block:
|
|
# Check if value is array-like - if so, use upsert_column logic
|
|
if isinstance(value, (pd.Series, np.ndarray)):
|
|
return self.upsert_column(name, value)
|
|
# Scalar value - use original fill_column logic
|
|
return self._table.assign(**{name: value})
|
|
|
|
def slice(self, start: int, end: int, copy: bool = False) -> "pandas.DataFrame":
|
|
view = self._table[start:end]
|
|
view.reset_index(drop=True, inplace=True)
|
|
if copy:
|
|
view = view.copy(deep=True)
|
|
return view
|
|
|
|
def take(self, indices: List[int]) -> "pandas.DataFrame":
|
|
table = self._table.take(indices)
|
|
table.reset_index(drop=True, inplace=True)
|
|
return table
|
|
|
|
def drop(self, columns: List[str]) -> Block:
|
|
return self._table.drop(columns, axis="columns")
|
|
|
|
def select(self, columns: List[str]) -> "pandas.DataFrame":
|
|
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"Pandas blocks, but got: {columns}."
|
|
)
|
|
return self._table[columns]
|
|
|
|
def rename_columns(self, columns_rename: Dict[str, str]) -> "pandas.DataFrame":
|
|
return self._table.rename(columns=columns_rename, inplace=False, copy=False)
|
|
|
|
def upsert_column(
|
|
self, column_name: str, column_data: BlockColumn
|
|
) -> "pandas.DataFrame":
|
|
import pyarrow
|
|
|
|
if isinstance(column_data, (pyarrow.Array, pyarrow.ChunkedArray)):
|
|
column_data = column_data.to_pandas()
|
|
|
|
return self._table.assign(**{column_name: column_data})
|
|
|
|
def random_shuffle(self, random_seed: Optional[int]) -> "pandas.DataFrame":
|
|
table = self._table.sample(frac=1, random_state=random_seed)
|
|
table.reset_index(drop=True, inplace=True)
|
|
return table
|
|
|
|
def schema(self) -> PandasBlockSchema:
|
|
dtypes = self._table.dtypes
|
|
schema = PandasBlockSchema(
|
|
names=tuple(dtypes.index.tolist()),
|
|
types=tuple(dtypes.values.tolist()),
|
|
)
|
|
# Column names with non-str types of a pandas DataFrame is not
|
|
# supported by Ray Dataset.
|
|
if any(not isinstance(name, str) for name in schema.names):
|
|
raise ValueError(
|
|
"A Pandas DataFrame with column names of non-str types"
|
|
" is not supported by Ray Dataset. Column names of this"
|
|
f" DataFrame: {schema.names!r}."
|
|
)
|
|
return schema
|
|
|
|
def to_pandas(self) -> "pandas.DataFrame":
|
|
from ray.data.util.data_batch_conversion import _cast_tensor_columns_to_ndarrays
|
|
|
|
ctx = DataContext.get_current()
|
|
table = self._table
|
|
if ctx.enable_tensor_extension_casting:
|
|
table = _cast_tensor_columns_to_ndarrays(table)
|
|
return table
|
|
|
|
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.columns.tolist()
|
|
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.columns)
|
|
for column in columns:
|
|
if column not in column_names_set:
|
|
raise ValueError(
|
|
f"Cannot find column {column}, available columns: "
|
|
f"{self._table.columns.tolist()}"
|
|
)
|
|
|
|
arrays = []
|
|
for column in columns:
|
|
arrays.append(self._table[column].to_numpy())
|
|
|
|
if should_be_single_ndarray:
|
|
arrays = arrays[0]
|
|
else:
|
|
arrays = dict(zip(columns, arrays))
|
|
return arrays
|
|
|
|
def to_arrow(self) -> "pyarrow.Table":
|
|
import pyarrow as pa
|
|
|
|
from ray.data._internal.tensor_extensions.pandas import TensorDtype
|
|
|
|
# _from_pandas_safe handles object-dtype columns that pa.Table.from_pandas
|
|
# cannot convert (e.g. multi-dimensional numpy arrays, PIL images), because Arrow cannot handle them natively.
|
|
arrow_table = _from_pandas_safe(self._table)
|
|
|
|
# NOTE: Pandas by default coerces all-null column types (including None,
|
|
# NaN, etc) into "double" type by default, which is incorrect in a
|
|
# a lot of cases.
|
|
#
|
|
# To fix that, we traverse all the columns after conversion and
|
|
# replace all-null ones with the column of null-type that allows
|
|
# these columns to be properly combined with the same column
|
|
# containing non-null values and carrying appropriate type later.
|
|
null_coerced_columns = {}
|
|
|
|
for idx, col_name in enumerate(self._table.columns):
|
|
col = self._table[col_name]
|
|
|
|
# Skip coercing tensors to null-type to avoid type information loss
|
|
# See https://github.com/ray-project/ray/issues/59087 for context
|
|
if isinstance(col.dtype, (TensorDtype, pd.ArrowDtype)):
|
|
continue
|
|
|
|
if not col.notna().any():
|
|
# If there are only null-values, coerce column to Arrow's `NullType`
|
|
null_coerced_columns[(idx, col_name)] = pa.nulls(
|
|
len(col), type=pa.null()
|
|
)
|
|
|
|
# NOTE: We're updating columns in place to preserve any potential metadata
|
|
# set from conversion from original Pandas data-frame
|
|
for (idx, col_name), null_col in null_coerced_columns.items():
|
|
arrow_table = arrow_table.set_column(idx, col_name, null_col)
|
|
|
|
return arrow_table
|
|
|
|
def num_rows(self) -> int:
|
|
return self._table.shape[0]
|
|
|
|
def size_bytes(self) -> int:
|
|
from ray.data._internal.tensor_extensions.pandas import TensorArray
|
|
from ray.data.extensions import TensorArrayElement, TensorDtype
|
|
|
|
pd = lazy_import_pandas()
|
|
|
|
def get_deep_size(obj):
|
|
"""Calculates the memory size of objects,
|
|
including nested objects using an iterative approach."""
|
|
seen = set()
|
|
total_size = 0
|
|
objects = collections.deque([obj])
|
|
while objects:
|
|
current = objects.pop()
|
|
|
|
# Skip interning-eligible immutable objects
|
|
if isinstance(current, (str, bytes, int, float)):
|
|
size = sys.getsizeof(current)
|
|
total_size += size
|
|
continue
|
|
|
|
# Check if the object has been seen before
|
|
# i.e. a = np.ndarray([1,2,3]), b = [a,a]
|
|
# The patten above will have only one memory copy
|
|
if id(current) in seen:
|
|
continue
|
|
seen.add(id(current))
|
|
|
|
try:
|
|
size = sys.getsizeof(current)
|
|
except TypeError:
|
|
size = 0
|
|
total_size += size
|
|
|
|
# Handle specific cases
|
|
if isinstance(current, np.ndarray):
|
|
total_size += current.nbytes - size # Avoid double counting
|
|
elif isinstance(current, pd.DataFrame):
|
|
total_size += (
|
|
current.memory_usage(index=True, deep=True).sum() - size
|
|
)
|
|
elif isinstance(current, (list, tuple, set)):
|
|
objects.extend(current)
|
|
elif isinstance(current, dict):
|
|
objects.extend(current.keys())
|
|
objects.extend(current.values())
|
|
elif isinstance(current, TensorArrayElement):
|
|
objects.extend(current.to_numpy())
|
|
return total_size
|
|
|
|
# Get initial memory usage.
|
|
# No need for deep inspection here, as we will handle the str, object and
|
|
# extension columns separately.
|
|
memory_usage = self._table.memory_usage(index=True, deep=False)
|
|
|
|
# TensorDtype for ray.data._internal.tensor_extensions.pandas.TensorDtype
|
|
object_need_check = (TensorDtype,)
|
|
max_sample_count = _PANDAS_SIZE_BYTES_MAX_SAMPLE_COUNT
|
|
|
|
# Handle object columns separately
|
|
for column in self._table.columns:
|
|
# For str, object and extension dtypes, we calculate the size
|
|
# by sampling the data.
|
|
dtype = self._table[column].dtype
|
|
if (
|
|
is_string_dtype(dtype)
|
|
or is_object_dtype(dtype)
|
|
or isinstance(dtype, object_need_check)
|
|
):
|
|
total_size = len(self._table[column])
|
|
|
|
# Determine the sample size based on max_sample_count
|
|
sample_size = min(total_size, max_sample_count)
|
|
# Skip size calculation for empty columns
|
|
if sample_size == 0:
|
|
continue
|
|
if sample_size == total_size:
|
|
# Sampling the whole column: read values directly to avoid the
|
|
# permutation/copy overhead of .sample(). No randomness here, so
|
|
# this is trivially deterministic.
|
|
sampled_data = self._table[column].values
|
|
else:
|
|
# Use a fixed random_state so size_bytes() is deterministic
|
|
# across calls. Non-deterministic size estimation can cause
|
|
# streaming generator tasks to produce different block counts
|
|
# across replay attempts (e.g. lineage reconstruction), which
|
|
# surfaces as a silent hang or silent data loss downstream.
|
|
sampled_data = (
|
|
self._table[column].sample(n=sample_size, random_state=0).values
|
|
)
|
|
|
|
try:
|
|
if isinstance(sampled_data, TensorArray) and np.issubdtype(
|
|
sampled_data[0].numpy_dtype, np.number
|
|
):
|
|
column_memory_sample = sampled_data.nbytes
|
|
else:
|
|
vectorized_size_calc = np.vectorize(lambda x: get_deep_size(x))
|
|
column_memory_sample = np.sum(
|
|
vectorized_size_calc(sampled_data)
|
|
)
|
|
# Scale back to the full column size if we sampled
|
|
column_memory = column_memory_sample * (total_size / sample_size)
|
|
# Add the data memory usage on top of the index memory usage.
|
|
memory_usage[column] += int(column_memory)
|
|
except Exception as e:
|
|
# Handle or log the exception as needed
|
|
logger.warning(f"Error calculating size for column '{column}': {e}")
|
|
|
|
# Sum up total memory usage
|
|
total_memory_usage = memory_usage.sum()
|
|
|
|
return int(total_memory_usage)
|
|
|
|
def _zip(self, acc: BlockAccessor) -> "pandas.DataFrame":
|
|
r = self.to_pandas().copy(deep=False)
|
|
s = acc.to_pandas()
|
|
for col_name in s.columns:
|
|
col = s[col_name]
|
|
column_names = list(r.columns)
|
|
# Ensure the column names are unique after zip.
|
|
if col_name in column_names:
|
|
i = 1
|
|
new_name = col_name
|
|
while new_name in column_names:
|
|
new_name = "{}_{}".format(col_name, i)
|
|
i += 1
|
|
col_name = new_name
|
|
r[col_name] = col
|
|
return r
|
|
|
|
@staticmethod
|
|
def builder() -> PandasBlockBuilder:
|
|
return PandasBlockBuilder()
|
|
|
|
@staticmethod
|
|
def _empty_table() -> "pandas.DataFrame":
|
|
return PandasBlockBuilder._empty_table()
|
|
|
|
def _sample(self, n_samples: int, sort_key: "SortKey") -> "pandas.DataFrame":
|
|
return self._table[sort_key.get_columns()].sample(n_samples, ignore_index=True)
|
|
|
|
def sort(self, sort_key: "SortKey"):
|
|
assert (
|
|
sort_key.get_columns()
|
|
), f"Sorting columns couldn't be empty (got {sort_key.get_columns()})"
|
|
|
|
if self._table.shape[0] == 0:
|
|
return self._empty_table()
|
|
|
|
columns, ascending = sort_key.to_pandas_sort_args()
|
|
return self._table.sort_values(by=columns, ascending=ascending)
|
|
|
|
def sort_and_partition(
|
|
self, boundaries: List[T], sort_key: "SortKey"
|
|
) -> List[Block]:
|
|
table = self.sort(sort_key)
|
|
|
|
if table.shape[0] == 0:
|
|
# If the pyarrow table is empty we may not have schema
|
|
# so calling sort_indices() will raise an error.
|
|
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"]:
|
|
pd = lazy_import_pandas()
|
|
stats = BlockExecStats.builder()
|
|
blocks = [b for b in blocks if b.shape[0] > 0]
|
|
if len(blocks) == 0:
|
|
ret = PandasBlockAccessor._empty_table()
|
|
else:
|
|
# Handle blocks of different types.
|
|
blocks = TableBlockAccessor.normalize_block_types(blocks, BlockType.PANDAS)
|
|
ret = pd.concat(blocks, ignore_index=True)
|
|
columns, ascending = sort_key.to_pandas_sort_args()
|
|
ret = ret.sort_values(by=columns, ascending=ascending)
|
|
from ray.data.block import BlockMetadataWithSchema
|
|
|
|
return ret, BlockMetadataWithSchema.from_block(
|
|
ret, block_exec_stats=stats.build()
|
|
)
|
|
|
|
def block_type(self) -> BlockType:
|
|
return BlockType.PANDAS
|
|
|
|
def iter_rows(
|
|
self, public_row_format: bool
|
|
) -> Iterator[Union[Mapping, np.ndarray]]:
|
|
num_rows = self.num_rows()
|
|
for i in range(num_rows):
|
|
row = self._get_row(i)
|
|
if public_row_format:
|
|
yield row.as_pydict()
|
|
else:
|
|
yield row
|
|
|
|
def filter(self, predicate_expr: "Expr") -> "pandas.DataFrame":
|
|
"""Filter rows based on a predicate expression."""
|
|
if self._table.empty:
|
|
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 pandas boolean indexing
|
|
return self._table[mask]
|