chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,776 @@
|
||||
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]
|
||||
Reference in New Issue
Block a user