ray/python/ray/data/block.py at master · ray-project/ray

960 lines (760 loc) · 32.1 KB
import collections
import logging
import time
from dataclasses import dataclass, field, fields
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Callable,
    Iterator,
    Optional,
    Protocol,
    Sequence,
    TypeVar,
import numpy as np
import pyarrow as pa
from ray.data._internal.util import _check_pyarrow_version, _truncated_repr
from ray.data.context import DataContext
from ray.types import ObjectRef
from ray.util import log_once
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
    import cudf
    import pandas
    import pyarrow
    from ray.data._internal.block_builder import BlockBuilder
    from ray.data._internal.pandas_block import PandasBlockSchema
    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
    from ray.data.aggregate import AggregateFn
T = TypeVar("T", contravariant=True)
U = TypeVar("U", covariant=True)
KeyType = TypeVar("KeyType")
AggType = TypeVar("AggType")
# Represents a batch of records to be stored in the Ray object store.
# Block data can be accessed in a uniform way via ``BlockAccessors`` like`
# ``ArrowBlockAccessor``.
Block = Union["pyarrow.Table", "pandas.DataFrame"]
# Represents the schema of a block, which can be either a Python type or a
# pyarrow schema. This is used to describe the structure of the data in a block.
Schema = Union["PandasBlockSchema", "pyarrow.lib.Schema"]
# Represents a single column of the ``Block``
BlockColumn = Union[
    "pyarrow.ChunkedArray",
    "pyarrow.Array",
    "pandas.Series",
# Represents a single column of the ``Batch``
BatchColumn = Union[
    "pandas.Series", "np.ndarray", "pyarrow.Array", "pyarrow.ChunkedArray"
logger = logging.getLogger(__name__)
@DeveloperAPI
class BlockType(Enum):
    ARROW = "arrow"
    PANDAS = "pandas"
@DeveloperAPI
class BatchFormat(str, Enum):
    # NOTE: This is to maintain compatibility w/ existing APIs
    ARROW = "pyarrow"
    PANDAS = "pandas"
    NUMPY = "numpy"
    CUDF = "cudf"
# User-facing data batch type. This is the data type for data that is supplied to and
# returned from batch UDFs.
DataBatch = Union[
    "pyarrow.Table",
    "pandas.DataFrame",
    Dict[str, np.ndarray],
    "cudf.DataFrame",
# User-facing data column type. This is the data type for data that is supplied to and
# returned from column UDFs.
DataBatchColumn = Union[BlockColumn, np.ndarray]
# A class type that implements __call__.
CallableClass = type
class _CallableClassProtocol(Protocol[T, U]):
    def __call__(self, __arg: T) -> Union[U, Iterator[U]]:
        ...
# A user defined function passed to flat_map, map_batches, etc.
UserDefinedFunction = Union[
    Callable[[T], U],
    Callable[[T], Iterator[U]],
    type["_CallableClassProtocol"],
# A list of block references pending computation by a single task. For example,
# this may be the output of a task reading a file.
BlockPartition = List[Tuple[ObjectRef[Block], "BlockMetadata"]]
# The metadata that describes the output of a BlockPartition. This has the
# same type as the metadata that describes each block in the partition.
BlockPartitionMetadata = List["BlockMetadata"]
VALID_BATCH_FORMATS = ["pandas", "pyarrow", "numpy", "cudf", None]
DEFAULT_BATCH_FORMAT = "numpy"
def _is_cudf_dataframe(obj: Any) -> bool:
    """Check if the object is a cudf.DataFrame (lazy import).
    Checks ``sys.modules`` first to avoid importing cudf (which loads CUDA
    and ~1.5 GiB of RSS) when it hasn't been imported yet.  If cudf is not
    in ``sys.modules``, no object in the process can be a cudf DataFrame.
    if "cudf" not in sys.modules:
        return False
        import cudf
        return isinstance(obj, cudf.DataFrame)
    except ImportError:
        return False
def _is_empty_schema(schema: Optional[Schema]) -> bool:
    from ray.data._internal.pandas_block import PandasBlockSchema
    return schema is None or (
        not schema.names
        if isinstance(schema, PandasBlockSchema)
        else not schema  # pyarrow schema check
def _take_first_non_empty_schema(schemas: Iterator["Schema"]) -> Optional["Schema"]:
    """Return the first non-empty schema from an iterator of schemas.
        schemas: Iterator of schemas to check.
    Returns:
        The first non-empty schema, or None if all schemas are empty.
    for schema in schemas:
        if not _is_empty_schema(schema):
            return schema
    return None
def _apply_batch_format(given_batch_format: Optional[str]) -> Optional[str]:
    if given_batch_format == "default":
        given_batch_format = DEFAULT_BATCH_FORMAT
    if given_batch_format not in VALID_BATCH_FORMATS:
        raise ValueError(
            f"The given batch format {given_batch_format} isn't allowed (must be one of"
            f" {VALID_BATCH_FORMATS})."
    return given_batch_format
@DeveloperAPI
def to_stats(metas: List["BlockMetadata"]) -> List["BlockStats"]:
    return [m.to_stats() for m in metas]
@DeveloperAPI
@dataclass(frozen=True)
class TaskExecWorkerStats:
    """Task's execution stats reported from the executing worker"""
    # Total task's wall-clock time from start to finish (measured on the worker)
    task_wall_time_s: float
@DeveloperAPI
@dataclass(frozen=True)
class BlockExecStats:
    """Execution stats for a single output block produced by a task."""
    # Index of the task that produced this block, used to attribute rows
    # to individual tasks in per-task statistics.
    task_idx: Optional[int] = None
    # Ray node ID of the worker that produced this block.
    node_id: str = field(
        default_factory=lambda: ray.runtime_context.get_runtime_context().get_node_id()
    # Absolute wall-clock timestamp when block generation started.
    start_time_s: Optional[float] = None
    # Absolute wall-clock timestamp when block generation finished.
    end_time_s: Optional[float] = None
    # Total wall-clock duration of the block generation (computed as end_time_s - start_time_s).
    wall_time_s: Optional[float] = None
    # Time spent inside UDF while generating block.
    udf_time_s: Optional[float] = 0
    # Time spent serializing this block into a Ray object.
    block_ser_time_s: Optional[float] = None
    # Total CPU time consumed by the worker process during the task, across all threads.
    cpu_time_s: Optional[float] = None
    # Peak USS (Unique Set Size) memory in bytes observed while computing this block,
    # as estimated by the memory profiler.
    max_uss_bytes: int = 0
    @staticmethod
    def builder() -> "_BlockExecStatsBuilder":
        return _BlockExecStatsBuilder()
class _BlockExecStatsBuilder:
    """Helper class for building block stats.
    When this class is created, we record the start time. When build() is
    called, the time delta is saved as part of the stats.
    def __init__(self):
        self._start_time = time.perf_counter()
        self._start_cpu = time.process_time()
        self._end_time = None
        self._end_cpu = None
    def finish(self):
        """Capture timing now, to be used by a later build() call."""
        self._end_time = time.perf_counter()
        self._end_cpu = time.process_time()
    def build(self, **kwargs) -> "BlockExecStats":
        if self._end_time is None:
            self.finish()
        return BlockExecStats(
            start_time_s=self._start_time,
            end_time_s=self._end_time,
            wall_time_s=self._end_time - self._start_time,
            cpu_time_s=self._end_cpu - self._start_cpu,
            **kwargs,
@DeveloperAPI
@dataclass(frozen=True)
class BlockStats:
    """Statistics about the block produced"""
    # The number of rows contained in this block, or None.
    num_rows: Optional[int]
    # The approximate size in bytes of this block, or None.
    size_bytes: Optional[int]
    # Execution stats for this block.
    exec_stats: Optional[BlockExecStats]
    # Overall task execution stats (reported from the worker).
    task_exec_stats: Optional[TaskExecWorkerStats] = field(default=None)
    def __post_init__(self):
        if self.size_bytes is not None:
            # Require size_bytes to be int, ray.util.metrics objects
            # will not take other types like numpy.int64
            assert isinstance(self.size_bytes, int)
_BLOCK_STATS_FIELD_NAMES = {f.name for f in fields(BlockStats)}
@DeveloperAPI
@dataclass(frozen=True)
class BlockMetadata(BlockStats):
    """Metadata about the block."""
    # The pyarrow schema or types of the block elements, or None.
    # The list of file paths used to generate this block, or
    # the empty list if indeterminate.
    # Stored as a tuple for hash-ability.
    input_files: Optional[Tuple[str, ...]] = field(default=None)
    def __post_init__(self):
        super().__post_init__()
        if self.input_files is not None and not isinstance(self.input_files, tuple):
            object.__setattr__(self, "input_files", tuple(self.input_files))
    def to_stats(self):
        return BlockStats(
            **{key: self.__getattribute__(key) for key in _BLOCK_STATS_FIELD_NAMES}
@DeveloperAPI(stability="alpha")
@dataclass(frozen=True)
class BlockMetadataWithSchema(BlockMetadata):
    schema: Optional[Schema] = None
    @staticmethod
    def from_metadata(
        metadata: "BlockMetadata", schema: Optional["Schema"] = None
    ) -> "BlockMetadataWithSchema":
        return BlockMetadataWithSchema(
            num_rows=metadata.num_rows,
            size_bytes=metadata.size_bytes,
            exec_stats=metadata.exec_stats,
            task_exec_stats=metadata.task_exec_stats,
            input_files=metadata.input_files,
            schema=schema,
    @staticmethod
    def from_block(
        block: Block,
        block_exec_stats: Optional["BlockExecStats"] = None,
        task_exec_stats: Optional["TaskExecWorkerStats"] = None,
    ) -> "BlockMetadataWithSchema":
        accessor = BlockAccessor.for_block(block)
        return BlockMetadataWithSchema.from_metadata(
            metadata=accessor.get_metadata(
                block_exec_stats=block_exec_stats,
                task_exec_stats=task_exec_stats,
            schema=accessor.schema(),
    @property
    def metadata(self) -> BlockMetadata:
        return BlockMetadata(
            num_rows=self.num_rows,
            size_bytes=self.size_bytes,
            exec_stats=self.exec_stats,
            input_files=self.input_files,
            task_exec_stats=self.task_exec_stats,
    def __getstate__(self) -> Dict[str, Any]:
        state = {f.name: getattr(self, f.name) for f in fields(BlockMetadataWithSchema)}
        if isinstance(self.schema, pa.Schema):
            state["schema"] = self.schema.serialize().to_pybytes()
        else:
            state["schema"] = self.schema
        return state
    def __setstate__(self, state: Dict[str, Any]):
        schema_val: bytes | Schema | None = state["schema"]
        if isinstance(schema_val, (bytes, bytearray)):
            state["schema"] = pa.ipc.read_schema(pa.BufferReader(schema_val))
        self.__dict__.update(state)
@DeveloperAPI
class BlockAccessor:
    """Provides accessor methods for a specific block.
    Ideally, we wouldn't need a separate accessor classes for blocks. However,
    this is needed if we want to support storing ``pyarrow.Table`` directly
    as a top-level Ray object, without a wrapping class (issue #17186).
    def num_rows(self) -> int:
        """Return the number of rows contained in this block."""
        raise NotImplementedError
    def iter_rows(self, public_row_format: bool) -> Iterator[T]:
        """Iterate over the rows of this block.
        Args:
            public_row_format: Whether to cast rows into the public Dict row
                format (this incurs extra copy conversions).
        """
        raise NotImplementedError
    def slice(self, start: int, end: int, copy: bool = False) -> Block:
        """Return a slice of this block.
        Args:
            start: The starting index of the slice (inclusive).
            end: The ending index of the slice (exclusive).
            copy: Whether to perform a data copy for the slice.
        Returns:
            The sliced block result.
        """
        raise NotImplementedError
    def take(self, indices: List[int]) -> Block:
        """Return a new block containing the provided row indices.
        Args:
            indices: The row indices to return.
        Returns:
            A new block containing the provided row indices.
        """
        raise NotImplementedError
    def drop(self, columns: List[str]) -> Block:
        """Return a new block with the list of provided columns dropped"""
        raise NotImplementedError
    def select(self, columns: List[Optional[str]]) -> Block:
        """Return a new block containing the provided columns."""
        raise NotImplementedError
    def rename_columns(self, columns_rename: Dict[str, str]) -> Block:
        """Return the block reflecting the renamed columns."""
        raise NotImplementedError
    def upsert_column(self, column_name: str, column_data: BlockColumn) -> Block:
        """
        Upserts a column into the block. If the column already exists, it will be replaced.
        Args:
            column_name: The name of the column to upsert.
            column_data: The data to upsert into the column. (Arrow Array/ChunkedArray for Arrow blocks, Series or array-like for Pandas blocks)
        Returns:
            The updated block.
        """
        raise NotImplementedError()
    def random_shuffle(self, random_seed: Optional[int]) -> Block:
        """Randomly shuffle this block."""
        raise NotImplementedError
    def to_pandas(self) -> "pandas.DataFrame":
        """Convert this block into a Pandas dataframe."""
        raise NotImplementedError
    def to_numpy(
        self, columns: Optional[Union[str, List[str]]] = None
    ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
        """Convert this block (or columns of block) into a NumPy ndarray.
        Args:
            columns: Name of columns to convert, or None if converting all columns.
        """
        raise NotImplementedError
    def to_arrow(self) -> "pyarrow.Table":
        """Convert this block into an Arrow table."""
        raise NotImplementedError
    def to_block(self) -> Block:
        """Return the base block that this accessor wraps."""
        raise NotImplementedError
    def to_default(self) -> Block:
        """Return the default data format for this accessor."""
        return self.to_block()
    def to_batch_format(self, batch_format: Optional[str]) -> DataBatch:
        """Convert this block into the provided batch format.
        Args:
            batch_format: The batch format to convert this block to.
        Returns:
            This block formatted as the provided batch format.
        """
        if batch_format is None:
            return self.to_block()
        elif batch_format == "default" or batch_format == "native":
            return self.to_default()
        elif batch_format == "pandas":
            return self.to_pandas()
        elif batch_format == "pyarrow":
            return self.to_arrow()
        elif batch_format == "numpy":
            return self.to_numpy()
        elif batch_format == "cudf":
            return self.to_cudf()
        else:
            raise ValueError(
                f"The batch format must be one of {VALID_BATCH_FORMATS}, got: "
                f"{batch_format}"
    def size_bytes(self) -> int:
        """Return the approximate size in bytes of this block."""
        raise NotImplementedError
    def schema(self) -> Union[type, "pyarrow.lib.Schema"]:
        """Return the Python type or pyarrow schema of this block."""
        raise NotImplementedError
    def get_metadata(
        self,
        input_files: Optional[List[str]] = None,
        block_exec_stats: Optional[BlockExecStats] = None,
        task_exec_stats: Optional[TaskExecWorkerStats] = None,
    ) -> BlockMetadata:
        """Create a metadata object from this block."""
        return BlockMetadata(
            num_rows=self.num_rows(),
            size_bytes=self.size_bytes(),
            input_files=tuple(input_files) if input_files is not None else None,
            exec_stats=block_exec_stats,
            task_exec_stats=task_exec_stats,
    def zip(self, other: "Block") -> "Block":
        """Zip this block with another block of the same type and size."""
        raise NotImplementedError
    @staticmethod
    def builder() -> "BlockBuilder":
        """Create a builder for this block type."""
        raise NotImplementedError
    @classmethod
    def batch_to_block(
        cls,
        batch: DataBatch,
        block_type: Optional[BlockType] = None,
    ) -> Block:
        """Create a block from user-facing data formats."""
        import pandas
        if isinstance(batch, np.ndarray):
            raise ValueError(
                f"Error validating {_truncated_repr(batch)}: "
                "Standalone numpy arrays are not "
                "allowed in Ray 2.5. Return a dict of field -> array, "
                "e.g., `{'data': array}` instead of `array`."
        # Handle cudf.DataFrame before Mapping check, since cudf.DataFrame
        # implements the Mapping protocol. Use bulk GPU->CPU transfer via
        # to_arrow() instead of the slow column-by-column Mapping path.
        elif _is_cudf_dataframe(batch):
            return batch.to_arrow(preserve_index=False)
        elif isinstance(batch, pandas.DataFrame):
            if (block_type == BlockType.ARROW) or (
                block_type is None
                and DataContext.get_current().batch_to_block_arrow_format
                return cls.for_block(batch).to_arrow()
            return batch
        elif isinstance(batch, collections.abc.Mapping):
            if block_type is None or block_type == BlockType.ARROW:
                from ray.data._internal.tensor_extensions.arrow import (
                    ArrowConversionError,
                    return cls.batch_to_arrow_block(batch)
                except ArrowConversionError as e:
                    if log_once("_fallback_to_pandas_block_warning"):
                        logger.debug(
                            f"Failed to convert batch to Arrow due to: {e}; "
                            f"falling back to Pandas block"
                    if block_type is None:
                        return cls.batch_to_pandas_block(batch)
            else:
                assert block_type == BlockType.PANDAS
                return cls.batch_to_pandas_block(batch)
        return batch
    @classmethod
    def batch_to_arrow_block(cls, batch: Dict[str, Any]) -> Block:
        """Create an Arrow block from user-facing data formats."""
        from ray.data._internal.arrow_block import ArrowBlockBuilder
        return ArrowBlockBuilder._table_from_pydict(batch)
    @classmethod
    def batch_to_pandas_block(cls, batch: Dict[str, Any]) -> Block:
        """Create a Pandas block from user-facing data formats."""
        from ray.data._internal.pandas_block import PandasBlockBuilder
        return PandasBlockBuilder._table_from_pydict(batch)
    @staticmethod
    def for_block(block: Block) -> "BlockAccessor[T]":
        """Create a block accessor for the given block."""
        _check_pyarrow_version()
        import pandas
        import pyarrow
        if isinstance(block, (pyarrow.Table, pyarrow.RecordBatch)):
            from ray.data._internal.arrow_block import ArrowBlockAccessor
            return ArrowBlockAccessor(block)
        elif isinstance(block, pandas.DataFrame):
            from ray.data._internal.pandas_block import PandasBlockAccessor
            return PandasBlockAccessor(block)
        elif isinstance(block, bytes):
            from ray.data._internal.arrow_block import ArrowBlockAccessor
            return ArrowBlockAccessor.from_bytes(block)
        elif isinstance(block, list):
            raise ValueError(
                f"Error validating {_truncated_repr(block)}: "
                "Standalone Python objects are not "
                "allowed in Ray 2.5. To use Python objects in a dataset, "
                "wrap them in a dict of numpy arrays, e.g., "
                "return `{'item': batch}` instead of just `batch`."
        else:
            raise TypeError("Not a block type: {} ({})".format(block, type(block)))
    def sample(self, n_samples: int, sort_key: "SortKey") -> "Block":
        """Return a random sample of items from this block."""
        raise NotImplementedError
    def count(self, on: str, ignore_nulls: bool = False) -> Optional[U]:
        """Returns a count of the distinct values in the provided column"""
        raise NotImplementedError
    def sum(self, on: str, ignore_nulls: bool) -> Optional[U]:
        """Returns a sum of the values in the provided column"""
        raise NotImplementedError
    def min(self, on: str, ignore_nulls: bool) -> Optional[U]:
        """Returns a min of the values in the provided column"""
        raise NotImplementedError
    def max(self, on: str, ignore_nulls: bool) -> Optional[U]:
        """Returns a max of the values in the provided column"""
        raise NotImplementedError
    def mean(self, on: str, ignore_nulls: bool) -> Optional[U]:
        """Returns a mean of the values in the provided column"""
        raise NotImplementedError
    def sum_of_squared_diffs_from_mean(
        self,
        on: str,
        ignore_nulls: bool,
        mean: Optional[U] = None,
    ) -> Optional[U]:
        """Returns a sum of diffs (from mean) squared for the provided column"""
        raise NotImplementedError
    def sort(self, sort_key: "SortKey") -> "Block":
        """Returns new block sorted according to provided `sort_key`"""
        raise NotImplementedError
    def sort_and_partition(
        self, boundaries: List[T], sort_key: "SortKey"
    ) -> List["Block"]:
        """Return a list of sorted partitions of this block."""
        raise NotImplementedError
    def _aggregate(self, key: "SortKey", aggs: Tuple["AggregateFn"]) -> Block:
        """Combine rows with the same key into an accumulator."""
        raise NotImplementedError
    @staticmethod
    def merge_sorted_blocks(
        blocks: List["Block"], sort_key: "SortKey"
    ) -> Tuple[Block, BlockMetadataWithSchema]:
        """Return a sorted block by merging a list of sorted blocks."""
        raise NotImplementedError
    @staticmethod
    def _combine_aggregated_blocks(
        blocks: List[Block],
        sort_key: "SortKey",
        aggs: Tuple["AggregateFn"],
        finalize: bool = True,
    ) -> Tuple[Block, BlockMetadataWithSchema]:
        """Aggregate partially combined and sorted blocks."""
        raise NotImplementedError
    def _find_partitions_sorted(
        self,
        boundaries: List[Tuple[Any]],
        sort_key: "SortKey",
    ) -> List[Block]:
        """NOTE: PLEASE READ CAREFULLY
        Returns dataset partitioned using list of boundaries
        This method requires that
            - Block being sorted (according to `sort_key`)
            - Boundaries is a sorted list of tuples
        """
        raise NotImplementedError
    def block_type(self) -> BlockType:
        """Return the block type of this block."""
        raise NotImplementedError
    def _get_group_boundaries_sorted(self, keys: List[str]) -> np.ndarray:
        """
        NOTE: THIS METHOD ASSUMES THAT PROVIDED BLOCK IS ALREADY SORTED
        Compute boundaries of the groups within a block based on provided
        key (a column or a list of columns)
        NOTE: In each column, NaNs/None are considered to be the same group.
        Args:
            block: sorted block for which grouping of rows will be determined
                    based on provided key
            keys: list of columns determining the key for every row based on
                    which the block will be grouped
        Returns:
            A list of starting indices of each group and an end index of the last
            group, i.e., there are ``num_groups + 1`` entries and the first and last
            entries are 0 and ``len(array)`` respectively.
        """
        if self.num_rows() == 0:
            return np.array([], dtype=np.int32)
        elif not keys:
            # If no keys are specified, whole block is considered a single group
            return np.array([0, self.num_rows()])
        # Convert key columns to Numpy (to perform vectorized
        # ops on them)
        projected_block = self.to_numpy(keys)
        return _get_group_boundaries_sorted_numpy(list(projected_block.values()))
    def _iter_groups_sorted(
        self, sort_key: "SortKey"
    ) -> Iterator[Tuple[Sequence[KeyType], Block]]:
        """
        NOTE: THIS METHOD ASSUMES THE BLOCK BEING SORTED
        Creates an iterator over (zero-copy) blocks of rows grouped by
        provided key(s).
        """
        key_col_names: List[str] = sort_key.get_columns()
        if not key_col_names:
            # Global aggregation consists of a single "group", so we short-circuit.
            yield tuple(), self.to_block()
            return
        boundaries = self._get_group_boundaries_sorted(key_col_names)
        for start, end in zip(boundaries[:-1], boundaries[1:]):
            # Fetch tuple of key values from the first row
            row = self._get_row(start)
            yield row[key_col_names], self.slice(start, end, copy=False)
@DeveloperAPI(stability="beta")
class BlockColumnAccessor:
    """Provides vendor-neutral interface to apply common operations
    to block's (Pandas/Arrow) columns"""
    def __init__(self, col: BlockColumn):
        self._column = col
    def count(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
        """Returns a count of the distinct values in the column"""
        raise NotImplementedError()
    def sum(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
        """Returns a sum of the values in the column"""
        return NotImplementedError()
    def min(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
        """Returns a min of the values in the column"""
        raise NotImplementedError()
    def max(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
        """Returns a max of the values in the column"""
        raise NotImplementedError()
    def mean(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
        """Returns a mean of the values in the column"""
        raise NotImplementedError()
    def quantile(
        self, *, q: float, ignore_nulls: bool, as_py: bool = True
    ) -> Optional[U]:
        """Returns requested quantile of the given column"""
        raise NotImplementedError()
    def unique(self) -> BlockColumn:
        """Returns new column holding only distinct values of the current one"""
        raise NotImplementedError()
    def value_counts(self) -> Dict[str, List]:
        raise NotImplementedError()
    def hash(self) -> BlockColumn:
        """
        Computes a 64-bit hash value for each row in the column.
        Provides a unified hashing method across supported backends.
        Handles complex types like lists or nested structures by producing a single hash per row.
        These hashes are useful for downstream operations such as deduplication, grouping, or partitioning.
        Internally, Polars is used to compute row-level hashes even when the original column
        is backed by Pandas or PyArrow.
        Returns:
            A column of 64-bit integer hashes, returned in the same format as the
            underlying backend (e.g., Pandas Series or PyArrow Array).
        """
        raise NotImplementedError()
    def flatten(self) -> BlockColumn:
        """Flattens nested lists merging them into top-level container"""
        raise NotImplementedError()
    def dropna(self) -> BlockColumn:
        raise NotImplementedError()
    def is_composed_of_lists(self) -> bool:
        """
        Checks whether the column is composed of list-like elements.
        Returns:
            True if the column is made up of list-like values; False otherwise.
        """
        raise NotImplementedError()
    def sum_of_squared_diffs_from_mean(
        self,
        ignore_nulls: bool,
        mean: Optional[U] = None,
        as_py: bool = True,
    ) -> Optional[U]:
        """Returns a sum of diffs (from mean) squared for the column"""
        raise NotImplementedError()
    def to_pylist(self) -> List[Any]:
        """Converts block column to a list of Python native objects"""
        raise NotImplementedError()
    def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray:
        """Converts underlying column to Numpy"""
        raise NotImplementedError()
    def _to_arrow_compatible_container(self) -> Union[List[Any], "pyarrow.Array"]:
        """Converts block column into a representation compatible with Arrow"""
        raise NotImplementedError()
    @staticmethod
    def for_column(col: BlockColumn) -> "BlockColumnAccessor":
        """Create a column accessor for the given column"""
        _check_pyarrow_version()
        import pandas as pd
        if isinstance(col, pa.Array) or isinstance(col, pa.ChunkedArray):
            from ray.data._internal.arrow_block import ArrowBlockColumnAccessor
            return ArrowBlockColumnAccessor(col)
        elif isinstance(col, pd.Series):
            from ray.data._internal.pandas_block import PandasBlockColumnAccessor
            return PandasBlockColumnAccessor(col)
        else:
            raise TypeError(
                f"Expected either a pandas.Series or pyarrow.Array "
                f"(ChunkedArray) (got {type(col)})"
def _get_group_boundaries_sorted_numpy(columns: list[np.ndarray]) -> np.ndarray:
    # There are 3 categories: general, numerics with NaN, and categorical with None.
    # We only needed to check the last element for NaNs/None, as they are assumed to
    # be sorted.
    general_arrays = []
    num_arrays_with_nan = []
    cat_arrays_with_none = []
    for arr in columns:
        if np.issubdtype(arr.dtype, np.number) and np.isnan(arr[-1]):
            num_arrays_with_nan.append(arr)
        elif not np.issubdtype(arr.dtype, np.number) and arr[-1] is None:
            cat_arrays_with_none.append(arr)
        else:
            general_arrays.append(arr)
    # Compute the difference between each pair of elements. Handle the cases
    # where neighboring elements are both NaN or None. Output as a list of
    # boolean arrays.
    diffs = []
    if len(general_arrays) > 0:
        diffs.append(
            np.vstack([arr[1:] != arr[:-1] for arr in general_arrays]).any(axis=0)
    if len(num_arrays_with_nan) > 0:
        # Two neighboring numeric elements belong to the same group when they are
        # 1) both finite and equal
        # or 2) both np.nan
        diffs.append(
            np.vstack(
                    (arr[1:] != arr[:-1])
                    & (np.isfinite(arr[1:]) | np.isfinite(arr[:-1]))
                    for arr in num_arrays_with_nan
            ).any(axis=0)
    if len(cat_arrays_with_none) > 0:
        # Two neighboring str/object elements belong to the same group when they are
        # 1) both finite and equal
        # or 2) both None
        diffs.append(
            np.vstack(
                    (arr[1:] != arr[:-1])
                    & ~(np.equal(arr[1:], None) & np.equal(arr[:-1], None))
                    for arr in cat_arrays_with_none
            ).any(axis=0)
    # A series of vectorized operations to compute the boundaries:
    # - column_stack: stack the bool arrays into a single 2D bool array
    # - any() and nonzero(): find the indices where any of the column diffs are True
    # - add 1 to get the index of the first element of the next group
    # - hstack(): include the 0 and last indices to the boundaries
    boundaries = np.hstack(
            (np.column_stack(diffs).any(axis=1).nonzero()[0] + 1),
            [len(columns[0])],
    ).astype(int)
    return boundaries
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

block.py

Latest commit

History

block.py

File metadata and controls
