wehub-resource-sync
591 lines
21 KiB
Python
591 lines
21 KiB
Python
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
|