chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,590 @@
|
||||
import collections
|
||||
import heapq
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
from ray._common.utils import env_integer
|
||||
from ray.data._internal.block_builder import BlockBuilder
|
||||
from ray.data._internal.size_estimator import SizeEstimator
|
||||
from ray.data._internal.util import (
|
||||
NULL_SENTINEL,
|
||||
find_partition_index,
|
||||
is_nan,
|
||||
keys_equal,
|
||||
)
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockAccessor,
|
||||
BlockColumnAccessor,
|
||||
BlockExecStats,
|
||||
BlockMetadataWithSchema,
|
||||
BlockType,
|
||||
KeyType,
|
||||
U,
|
||||
)
|
||||
from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
||||
from ray.data.aggregate import AggregateFn
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# The max size of Python tuples to buffer before compacting them into a
|
||||
# table in the BlockBuilder.
|
||||
MAX_UNCOMPACTED_SIZE_BYTES = env_integer(
|
||||
"RAY_DATA_MAX_UNCOMPACTED_SIZE_BYTES", DEFAULT_TARGET_MAX_BLOCK_SIZE
|
||||
)
|
||||
|
||||
|
||||
class TableBlockBuilder(BlockBuilder):
|
||||
def __init__(self, block_type):
|
||||
# The set of uncompacted Python values buffered.
|
||||
self._columns = collections.defaultdict(list)
|
||||
# The set of compacted tables we have built so far.
|
||||
self._tables: List[Any] = []
|
||||
# Cursor into tables indicating up to which table we've accumulated table sizes.
|
||||
# This is used to defer table size calculation, which can be expensive for e.g.
|
||||
# Pandas DataFrames.
|
||||
# This cursor points to the first table for which we haven't accumulated a table
|
||||
# size.
|
||||
self._tables_size_cursor = 0
|
||||
# Accumulated table sizes, up to the table in _tables pointed to by
|
||||
# _tables_size_cursor.
|
||||
self._tables_size_bytes = 0
|
||||
# Size estimator for un-compacted table values.
|
||||
self._uncompacted_size = SizeEstimator()
|
||||
self._num_rows = 0
|
||||
self._num_uncompacted_rows = 0
|
||||
self._num_compactions = 0
|
||||
self._block_type = block_type
|
||||
|
||||
def add(self, item: Union[dict, Mapping]) -> None:
|
||||
if hasattr(item, "as_pydict"):
|
||||
item = item.as_pydict()
|
||||
if not isinstance(item, collections.abc.Mapping):
|
||||
raise ValueError(
|
||||
"Returned elements of an TableBlock must be of type `dict`, "
|
||||
"got {} (type {}).".format(item, type(item))
|
||||
)
|
||||
|
||||
# Fill in missing columns with None.
|
||||
for column_name in item:
|
||||
if column_name not in self._columns:
|
||||
self._columns[column_name] = [None] * self._num_uncompacted_rows
|
||||
|
||||
for column_name in self._columns:
|
||||
value = item.get(column_name)
|
||||
self._columns[column_name].append(value)
|
||||
|
||||
self._num_rows += 1
|
||||
self._num_uncompacted_rows += 1
|
||||
self._compact_if_needed()
|
||||
self._uncompacted_size.add(item)
|
||||
|
||||
def add_block(self, block: Any) -> None:
|
||||
if not isinstance(block, self._block_type):
|
||||
raise TypeError(
|
||||
f"Got a block of type {type(block)}, expected {self._block_type}."
|
||||
"If you are mapping a function, ensure it returns an "
|
||||
"object with the expected type. Block:\n"
|
||||
f"{block}"
|
||||
)
|
||||
accessor = BlockAccessor.for_block(block)
|
||||
self._tables.append(block)
|
||||
self._num_rows += accessor.num_rows()
|
||||
|
||||
@staticmethod
|
||||
def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _combine_tables(tables: List[Block]) -> Block:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _empty_table() -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _concat_would_copy() -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def will_build_yield_copy(self) -> bool:
|
||||
if self._columns:
|
||||
# Building a table from a dict of list columns always creates a copy.
|
||||
return True
|
||||
return self._concat_would_copy() and len(self._tables) > 1
|
||||
|
||||
def build(self) -> Block:
|
||||
# Preserve insertion order: previously-compacted tables (older) first,
|
||||
# then any rows added since the last compaction (newest) last.
|
||||
tables = list(self._tables)
|
||||
if self._columns:
|
||||
tables.append(self._table_from_pydict(self._columns))
|
||||
|
||||
if len(tables) == 0:
|
||||
return self._empty_table()
|
||||
else:
|
||||
return self._combine_tables(tables)
|
||||
|
||||
def num_rows(self) -> int:
|
||||
return self._num_rows
|
||||
|
||||
def num_blocks(self) -> int:
|
||||
return len(self._tables)
|
||||
|
||||
def get_estimated_memory_usage(self) -> int:
|
||||
if self._num_rows == 0:
|
||||
return 0
|
||||
for table in self._tables[self._tables_size_cursor :]:
|
||||
self._tables_size_bytes += BlockAccessor.for_block(table).size_bytes()
|
||||
self._tables_size_cursor = len(self._tables)
|
||||
return self._tables_size_bytes + self._uncompacted_size.size_bytes()
|
||||
|
||||
def _compact_if_needed(self) -> None:
|
||||
assert self._columns
|
||||
if self._uncompacted_size.size_bytes() < MAX_UNCOMPACTED_SIZE_BYTES:
|
||||
return
|
||||
block = self._table_from_pydict(self._columns)
|
||||
self.add_block(block)
|
||||
self._uncompacted_size = SizeEstimator()
|
||||
self._columns.clear()
|
||||
self._num_compactions += 1
|
||||
self._num_uncompacted_rows = 0
|
||||
|
||||
|
||||
class TableBlockAccessor(BlockAccessor):
|
||||
def __init__(self, table: Any):
|
||||
self._table = table
|
||||
|
||||
@staticmethod
|
||||
def _munge_conflict(name, count):
|
||||
return f"{name}_{count + 1}"
|
||||
|
||||
def to_default(self) -> Block:
|
||||
# Always promote Arrow blocks to pandas for consistency, since
|
||||
# we lazily convert pandas->Arrow internally for efficiency.
|
||||
default = self.to_pandas()
|
||||
|
||||
return default
|
||||
|
||||
def to_cudf(self) -> Any:
|
||||
"""Convert this block to a cudf.DataFrame (requires cudf to be installed)."""
|
||||
from ray.data.util.data_batch_conversion import _lazy_import_cudf
|
||||
|
||||
cudf = _lazy_import_cudf()
|
||||
if cudf is None:
|
||||
raise ValueError(
|
||||
"Attempted to convert data to cuDF DataFrame but cuDF "
|
||||
"is not installed. Please do `pip install cudf-cu12` to "
|
||||
"install cuDF (GPU required)."
|
||||
)
|
||||
|
||||
return cudf.DataFrame.from_arrow(self.to_arrow())
|
||||
|
||||
def column_names(self) -> List[str]:
|
||||
raise NotImplementedError
|
||||
|
||||
def fill_column(self, name: str, value: Any) -> Block:
|
||||
raise NotImplementedError
|
||||
|
||||
def to_block(self) -> Block:
|
||||
return self._table
|
||||
|
||||
def _zip(self, acc: BlockAccessor) -> "Block":
|
||||
raise NotImplementedError
|
||||
|
||||
def zip(self, other: "Block") -> "Block":
|
||||
acc = BlockAccessor.for_block(other)
|
||||
if not isinstance(acc, type(self)):
|
||||
if isinstance(self, TableBlockAccessor) and isinstance(
|
||||
acc, TableBlockAccessor
|
||||
):
|
||||
# If block types are different, but still both of TableBlock type, try
|
||||
# converting both to default block type before zipping.
|
||||
self_norm, other_norm = TableBlockAccessor.normalize_block_types(
|
||||
[self._table, other],
|
||||
)
|
||||
return BlockAccessor.for_block(self_norm).zip(other_norm)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot zip {} with block of type {}".format(
|
||||
type(self), type(other)
|
||||
)
|
||||
)
|
||||
if acc.num_rows() != self.num_rows():
|
||||
raise ValueError(
|
||||
"Cannot zip self (length {}) with block of length {}".format(
|
||||
self.num_rows(), acc.num_rows()
|
||||
)
|
||||
)
|
||||
return self._zip(acc)
|
||||
|
||||
@staticmethod
|
||||
def _empty_table() -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
def _sample(self, n_samples: int, sort_key: "SortKey") -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
def sample(self, n_samples: int, sort_key: "SortKey") -> Any:
|
||||
if sort_key is None or callable(sort_key):
|
||||
raise NotImplementedError(
|
||||
f"Table sort key must be a column name, was: {sort_key}"
|
||||
)
|
||||
if self.num_rows() == 0:
|
||||
# If the pyarrow table is empty we may not have schema
|
||||
# so calling table.select() will raise an error.
|
||||
return self._empty_table()
|
||||
k = min(n_samples, self.num_rows())
|
||||
return self._sample(k, sort_key)
|
||||
|
||||
def count(self, on: str, ignore_nulls: bool = False) -> Optional[U]:
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.count(ignore_nulls=ignore_nulls)
|
||||
|
||||
def sum(self, on: str, ignore_nulls: bool) -> Optional[U]:
|
||||
self._validate_column(on)
|
||||
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.sum(ignore_nulls=ignore_nulls)
|
||||
|
||||
def min(self, on: str, ignore_nulls: bool) -> Optional[U]:
|
||||
self._validate_column(on)
|
||||
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.min(ignore_nulls=ignore_nulls)
|
||||
|
||||
def max(self, on: str, ignore_nulls: bool) -> Optional[U]:
|
||||
self._validate_column(on)
|
||||
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.max(ignore_nulls=ignore_nulls)
|
||||
|
||||
def mean(self, on: str, ignore_nulls: bool) -> Optional[U]:
|
||||
self._validate_column(on)
|
||||
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.mean(ignore_nulls=ignore_nulls)
|
||||
|
||||
def sum_of_squared_diffs_from_mean(
|
||||
self,
|
||||
on: str,
|
||||
ignore_nulls: bool,
|
||||
mean: Optional[U] = None,
|
||||
) -> Optional[U]:
|
||||
self._validate_column(on)
|
||||
|
||||
accessor = BlockColumnAccessor.for_column(self._table[on])
|
||||
return accessor.sum_of_squared_diffs_from_mean(ignore_nulls=ignore_nulls)
|
||||
|
||||
def _validate_column(self, col: str):
|
||||
if col is None:
|
||||
raise ValueError(f"Provided `on` value has to be non-null (got '{col}')")
|
||||
elif col not in self.column_names():
|
||||
raise ValueError(
|
||||
f"Referencing column '{col}' not present in the schema: {self.schema()}"
|
||||
)
|
||||
|
||||
def _aggregate(self, sort_key: "SortKey", aggs: Tuple["AggregateFn"]) -> Block:
|
||||
"""Applies provided aggregations to groups of rows with the same key.
|
||||
|
||||
This assumes the block is already sorted by key in ascending order.
|
||||
|
||||
Args:
|
||||
sort_key: A column name or list of column names.
|
||||
If this is ``None``, place all rows in a single group.
|
||||
|
||||
aggs: The aggregations to do.
|
||||
|
||||
Returns:
|
||||
A sorted block of [k, v_1, ..., v_n] columns where k is the groupby
|
||||
key and v_i is the partially combined accumulator for the ith given
|
||||
aggregation.
|
||||
If key is None then the k column is omitted.
|
||||
"""
|
||||
keys: List[str] = sort_key.get_columns()
|
||||
|
||||
def iter_groups() -> Iterator[Tuple[Sequence[KeyType], Block]]:
|
||||
"""Creates an iterator over zero-copy group views."""
|
||||
if not keys:
|
||||
# Global aggregation consists of a single "group", so we short-circuit.
|
||||
yield tuple(), self.to_block()
|
||||
return
|
||||
|
||||
yield from self._iter_groups_sorted(sort_key)
|
||||
|
||||
builder = self.builder()
|
||||
for group_keys, group_view in iter_groups():
|
||||
# Aggregate.
|
||||
init_vals = group_keys
|
||||
if len(group_keys) == 1:
|
||||
init_vals = group_keys[0]
|
||||
|
||||
accumulators = [agg.init(init_vals) for agg in aggs]
|
||||
for i in range(len(aggs)):
|
||||
accessor = BlockAccessor.for_block(group_view)
|
||||
# Skip empty blocks
|
||||
if accessor.num_rows() > 0:
|
||||
accumulators[i] = aggs[i].accumulate_block(
|
||||
accumulators[i], group_view
|
||||
)
|
||||
|
||||
# Build the row.
|
||||
row = {}
|
||||
if keys:
|
||||
for k, gk in zip(keys, group_keys):
|
||||
row[k] = gk
|
||||
|
||||
count = collections.defaultdict(int)
|
||||
for agg, accumulator in zip(aggs, accumulators):
|
||||
name = agg.name
|
||||
# Check for conflicts with existing aggregation name.
|
||||
if count[name] > 0:
|
||||
name = self._munge_conflict(name, count[name])
|
||||
count[name] += 1
|
||||
row[name] = accumulator
|
||||
|
||||
builder.add(row)
|
||||
|
||||
return builder.build()
|
||||
|
||||
@classmethod
|
||||
def _combine_aggregated_blocks(
|
||||
cls,
|
||||
blocks: List[Block],
|
||||
sort_key: "SortKey",
|
||||
aggs: Tuple["AggregateFn"],
|
||||
finalize: bool = True,
|
||||
) -> Tuple[Block, "BlockMetadataWithSchema"]:
|
||||
"""Combine previously aggregated blocks.
|
||||
|
||||
This assumes blocks are already sorted by key in ascending order,
|
||||
so we can do merge sort to get all the rows with the same key.
|
||||
|
||||
Args:
|
||||
blocks: A list of partially combined and sorted blocks.
|
||||
sort_key: The column name of key or None for global aggregation.
|
||||
aggs: The aggregations to do.
|
||||
finalize: Whether to finalize the aggregation. This is used as an
|
||||
optimization for cases where we repeatedly combine partially
|
||||
aggregated groups.
|
||||
|
||||
Returns:
|
||||
A block of [k, v_1, ..., v_n] columns and its metadata where k is
|
||||
the groupby key and v_i is the corresponding aggregation result for
|
||||
the ith given aggregation.
|
||||
If key is None then the k column is omitted.
|
||||
"""
|
||||
|
||||
# Handle blocks of different types.
|
||||
blocks = TableBlockAccessor.normalize_block_types(blocks)
|
||||
|
||||
stats = BlockExecStats.builder()
|
||||
keys = sort_key.get_columns()
|
||||
|
||||
def _key_fn(r):
|
||||
if keys:
|
||||
return tuple(r[keys])
|
||||
else:
|
||||
return (0,)
|
||||
|
||||
# Replace `None`s and `np.nan` with NULL_SENTINEL to make sure
|
||||
# we can order the elements (both of these are incomparable)
|
||||
def safe_key_fn(r):
|
||||
values = _key_fn(r)
|
||||
return tuple(
|
||||
[NULL_SENTINEL if v is None or is_nan(v) else v for v in values]
|
||||
)
|
||||
|
||||
iter = heapq.merge(
|
||||
*[
|
||||
BlockAccessor.for_block(block).iter_rows(public_row_format=False)
|
||||
for block in blocks
|
||||
],
|
||||
key=safe_key_fn,
|
||||
)
|
||||
|
||||
next_row = None
|
||||
builder = BlockAccessor.for_block(blocks[0]).builder()
|
||||
|
||||
while True:
|
||||
try:
|
||||
if next_row is None:
|
||||
next_row = next(iter)
|
||||
|
||||
next_keys = _key_fn(next_row)
|
||||
next_key_columns = keys
|
||||
|
||||
def gen():
|
||||
nonlocal iter
|
||||
nonlocal next_row
|
||||
while keys_equal(_key_fn(next_row), next_keys):
|
||||
yield next_row
|
||||
try:
|
||||
next_row = next(iter)
|
||||
except StopIteration:
|
||||
next_row = None
|
||||
break
|
||||
|
||||
# Merge.
|
||||
first = True
|
||||
accumulators = [None] * len(aggs)
|
||||
resolved_agg_names = [None] * len(aggs)
|
||||
for r in gen():
|
||||
if first:
|
||||
count = collections.defaultdict(int)
|
||||
for i in range(len(aggs)):
|
||||
name = aggs[i].name
|
||||
# Check for conflicts with existing aggregation
|
||||
# name.
|
||||
if count[name] > 0:
|
||||
name = TableBlockAccessor._munge_conflict(
|
||||
name, count[name]
|
||||
)
|
||||
count[name] += 1
|
||||
resolved_agg_names[i] = name
|
||||
accumulators[i] = r[name]
|
||||
first = False
|
||||
else:
|
||||
for i in range(len(aggs)):
|
||||
accumulators[i] = aggs[i].merge(
|
||||
accumulators[i],
|
||||
r[resolved_agg_names[i]],
|
||||
)
|
||||
# Build the row.
|
||||
row = {}
|
||||
if keys:
|
||||
for col_name, next_key in zip(next_key_columns, next_keys):
|
||||
row[col_name] = next_key
|
||||
|
||||
for agg, agg_name, accumulator in zip(
|
||||
aggs, resolved_agg_names, accumulators
|
||||
):
|
||||
if finalize:
|
||||
row[agg_name] = agg.finalize(accumulator)
|
||||
else:
|
||||
row[agg_name] = accumulator
|
||||
|
||||
builder.add(row)
|
||||
except StopIteration:
|
||||
break
|
||||
|
||||
ret = builder.build()
|
||||
return ret, BlockMetadataWithSchema.from_block(
|
||||
ret, block_exec_stats=stats.build()
|
||||
)
|
||||
|
||||
def _find_partitions_sorted(
|
||||
self,
|
||||
boundaries: List[Tuple[Any]],
|
||||
sort_key: "SortKey",
|
||||
):
|
||||
partitions = []
|
||||
|
||||
# For each boundary value, count the number of items that are less
|
||||
# than it. Since the block is sorted, these counts partition the items
|
||||
# such that boundaries[i] <= x < boundaries[i + 1] for each x in
|
||||
# partition[i]. If `descending` is true, `boundaries` would also be
|
||||
# in descending order and we only need to count the number of items
|
||||
# *greater than* the boundary value instead.
|
||||
bounds = [
|
||||
find_partition_index(self._table, boundary, sort_key)
|
||||
for boundary in boundaries
|
||||
]
|
||||
|
||||
last_idx = 0
|
||||
for idx in bounds:
|
||||
partitions.append(self._table[last_idx:idx])
|
||||
last_idx = idx
|
||||
partitions.append(self._table[last_idx:])
|
||||
return partitions
|
||||
|
||||
@classmethod
|
||||
def normalize_block_types(
|
||||
cls,
|
||||
blocks: List[Block],
|
||||
target_block_type: Optional[BlockType] = None,
|
||||
) -> List[Block]:
|
||||
"""Normalize input blocks to the specified `normalize_type`. If the blocks
|
||||
are already all of the same type, returns original blocks.
|
||||
|
||||
Args:
|
||||
blocks: A list of TableBlocks to be normalized.
|
||||
target_block_type: The type to normalize the blocks to. If None,
|
||||
Ray Data chooses a type to minimize the amount of data conversions.
|
||||
|
||||
Returns:
|
||||
A list of blocks of the same type.
|
||||
"""
|
||||
seen_types: Dict[BlockType, int] = collections.defaultdict(int)
|
||||
|
||||
for block in blocks:
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
if not isinstance(block_accessor, TableBlockAccessor):
|
||||
raise ValueError(
|
||||
"Block type normalization is only supported for TableBlock, "
|
||||
f"but received block of type: {type(block)}."
|
||||
)
|
||||
|
||||
seen_types[block_accessor.block_type()] += 1
|
||||
|
||||
# If there's just 1 block-type and it's matching target-type, short-circuit
|
||||
if len(seen_types) == 1 and (
|
||||
target_block_type is None or [target_block_type] == list(seen_types.keys())
|
||||
):
|
||||
return blocks
|
||||
|
||||
# Pick the most prevalent block-type
|
||||
if target_block_type is None:
|
||||
_, target_block_type = sorted(
|
||||
seen_types.items(),
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)[0]
|
||||
|
||||
results = [
|
||||
cls.try_convert_block_type(block, target_block_type) for block in blocks
|
||||
]
|
||||
|
||||
if any(not isinstance(block, type(results[0])) for block in results):
|
||||
raise ValueError(
|
||||
"Expected all blocks to be of the same type after normalization, but "
|
||||
f"got different types: {[type(b) for b in results]}. "
|
||||
"Try using blocks of the same type to avoid the issue "
|
||||
"with block normalization."
|
||||
)
|
||||
return results
|
||||
|
||||
@classmethod
|
||||
def try_convert_block_type(cls, block: Block, block_type: BlockType):
|
||||
if block_type == BlockType.ARROW:
|
||||
return BlockAccessor.for_block(block).to_arrow()
|
||||
elif block_type == BlockType.PANDAS:
|
||||
return BlockAccessor.for_block(block).to_pandas()
|
||||
else:
|
||||
return BlockAccessor.for_block(block).to_default()
|
||||
|
||||
def hstack(self, other_block: Block) -> Block:
|
||||
"""Combine this table with another table horizontally (column-wise).
|
||||
This will append the columns.
|
||||
|
||||
Args:
|
||||
other_block: The table to hstack side-by-side with.
|
||||
|
||||
Returns:
|
||||
A new table with columns from both tables combined.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user