wehub-resource-sync
484 lines
16 KiB
Python
484 lines
16 KiB
Python
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
from ray.data._internal.logical.interfaces import SourceOperator
|
|
from ray.data.block import Block, BlockAccessor, BlockMetadata
|
|
from ray.exceptions import RayError
|
|
from ray.types import ObjectRef
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data._internal.logical.interfaces.operator import Operator
|
|
from ray.data.dataset import Dataset, Schema
|
|
|
|
_DATASET_REPR_ELLIPSIS = "…" # Ellipsis marker for truncated cells/rows.
|
|
_DATASET_REPR_MAX_ROWS = 10 # Total preview row budget when materialized.
|
|
_DATASET_REPR_HEAD_ROWS = 5 # Number of head rows to show before the gap.
|
|
_DATASET_REPR_MAX_COLUMN_WIDTH = 40 # Max width per column cell in the table.
|
|
_DATASET_REPR_GET_TIMEOUT_S = 30.0 # Timeout for fetching preview blocks.
|
|
|
|
__all__ = [
|
|
"build_dataset_ascii_repr",
|
|
"build_dataset_summary_repr",
|
|
]
|
|
|
|
|
|
def build_dataset_ascii_repr(
|
|
dataset: "Dataset",
|
|
schema: "Schema",
|
|
is_materialized: bool,
|
|
) -> str:
|
|
"""Render the dataset as a multi-line tabular string."""
|
|
columns = list(schema.names)
|
|
if not columns:
|
|
return build_dataset_summary_repr(dataset)
|
|
|
|
num_rows = dataset._meta_count()
|
|
head_rows: List[List[str]] = []
|
|
tail_rows: List[List[str]] = []
|
|
if is_materialized:
|
|
try:
|
|
head_data, tail_data, _ = _collect_materialized_rows_for_repr(
|
|
dataset, num_rows
|
|
)
|
|
head_rows = _format_rows_for_repr(head_data, columns)
|
|
tail_rows = _format_rows_for_repr(tail_data, columns)
|
|
except RayError:
|
|
head_rows = []
|
|
tail_rows = []
|
|
|
|
return _build_dataset_ascii_repr_from_rows(
|
|
schema=schema,
|
|
num_rows=num_rows,
|
|
dataset_name=dataset.name,
|
|
is_materialized=is_materialized,
|
|
head_rows=head_rows,
|
|
tail_rows=tail_rows,
|
|
)
|
|
|
|
|
|
def _build_dataset_ascii_repr_from_rows(
|
|
*,
|
|
schema: "Schema",
|
|
num_rows: Optional[int],
|
|
dataset_name: Optional[str],
|
|
is_materialized: bool,
|
|
head_rows: List[List[str]],
|
|
tail_rows: List[List[str]],
|
|
) -> str:
|
|
"""Render the dataset repr given schema metadata and preview rows."""
|
|
columns = list(schema.names)
|
|
num_cols = len(columns)
|
|
shape_line = f"shape: ({num_rows if num_rows is not None else '?'}, {num_cols})"
|
|
|
|
# Build header rows from schema.
|
|
dtype_strings = [_repr_format_dtype(t) for t in schema.types]
|
|
column_headers = [
|
|
_truncate_to_cell_width(str(col), _DATASET_REPR_MAX_COLUMN_WIDTH)
|
|
for col in columns
|
|
]
|
|
dtype_headers = [
|
|
_truncate_to_cell_width(dtype, _DATASET_REPR_MAX_COLUMN_WIDTH)
|
|
for dtype in dtype_strings
|
|
]
|
|
separator_row = ["---"] * len(columns)
|
|
|
|
# Assemble rows, including an ellipsis gap if needed.
|
|
show_gap = bool(head_rows) and bool(tail_rows)
|
|
display_rows: List[List[str]] = []
|
|
display_rows.extend(head_rows)
|
|
if show_gap:
|
|
display_rows.append([_DATASET_REPR_ELLIPSIS] * len(columns))
|
|
display_rows.extend(tail_rows)
|
|
|
|
# Render the table with computed column widths.
|
|
column_widths = _compute_column_widths(
|
|
column_headers, dtype_headers, separator_row, display_rows
|
|
)
|
|
|
|
table_lines = _render_table_lines(
|
|
column_headers,
|
|
dtype_headers,
|
|
separator_row,
|
|
display_rows,
|
|
column_widths,
|
|
)
|
|
|
|
# Append a summary line describing row coverage.
|
|
num_rows_shown = len(head_rows) + len(tail_rows)
|
|
summary_line = (
|
|
f"(Showing {num_rows_shown} of {num_rows} rows)"
|
|
if is_materialized
|
|
else "(Dataset isn't materialized)"
|
|
)
|
|
if is_materialized and num_rows is None:
|
|
summary_line = f"(Showing {num_rows_shown} of ? rows)"
|
|
|
|
components = []
|
|
if dataset_name is not None:
|
|
components.append(f"name: {dataset_name}")
|
|
components.extend([shape_line, "\n".join(table_lines), summary_line])
|
|
return "\n".join(components)
|
|
|
|
|
|
def _repr_format_dtype(dtype: object) -> str:
|
|
"""Format a dtype into a compact string for the schema row.
|
|
|
|
Dtypes may come from PyArrow, pandas/NumPy, or be plain Python types.
|
|
"""
|
|
if isinstance(dtype, type):
|
|
return dtype.__name__
|
|
name = getattr(dtype, "name", None)
|
|
if isinstance(name, str):
|
|
return name
|
|
return str(dtype)
|
|
|
|
|
|
def _collect_materialized_rows_for_repr(
|
|
dataset: "Dataset",
|
|
num_rows: Optional[int],
|
|
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], bool]:
|
|
"""Collect head/tail rows for preview and whether to show a gap row."""
|
|
block_entries: List[Tuple[ObjectRef, BlockMetadata]] = []
|
|
for ref_bundle in dataset.iter_internal_ref_bundles():
|
|
block_entries.extend(zip(ref_bundle.block_refs, ref_bundle.metadata))
|
|
|
|
if not block_entries:
|
|
return [], [], False
|
|
|
|
# Compute how many head/tail rows to show within the preview budget.
|
|
head_row_limit, tail_row_limit = _determine_preview_row_targets(num_rows)
|
|
block_cache: Dict[ObjectRef, Block] = {}
|
|
|
|
def _resolve_block(block_ref: ObjectRef) -> Block:
|
|
if block_ref not in block_cache:
|
|
block_cache[block_ref] = ray.get(
|
|
block_ref, timeout=_DATASET_REPR_GET_TIMEOUT_S
|
|
)
|
|
return block_cache[block_ref]
|
|
|
|
head_rows: List[Dict[str, Any]] = []
|
|
head_remaining = head_row_limit
|
|
for block_ref, _ in block_entries:
|
|
if head_remaining <= 0:
|
|
break
|
|
block = _resolve_block(block_ref)
|
|
accessor = BlockAccessor.for_block(block)
|
|
for row in accessor.iter_rows(public_row_format=True):
|
|
head_rows.append(row)
|
|
head_remaining -= 1
|
|
if head_remaining <= 0:
|
|
break
|
|
|
|
tail_rows: List[Dict[str, Any]] = []
|
|
tail_remaining = tail_row_limit
|
|
tail_parts: List[List[Dict[str, Any]]] = []
|
|
if tail_remaining > 0:
|
|
for block_ref, metadata in reversed(block_entries):
|
|
if tail_remaining <= 0:
|
|
break
|
|
block = _resolve_block(block_ref)
|
|
accessor = BlockAccessor.for_block(block)
|
|
total_rows = metadata.num_rows
|
|
if total_rows is None:
|
|
total_rows = accessor.num_rows()
|
|
if total_rows == 0:
|
|
continue
|
|
start = max(0, total_rows - tail_remaining)
|
|
sliced_block = accessor.slice(start, total_rows, copy=False)
|
|
slice_accessor = BlockAccessor.for_block(sliced_block)
|
|
block_rows = list(slice_accessor.iter_rows(public_row_format=True))
|
|
tail_parts.append(block_rows)
|
|
tail_remaining -= len(block_rows)
|
|
if tail_remaining <= 0:
|
|
break
|
|
|
|
for part in reversed(tail_parts):
|
|
tail_rows.extend(part)
|
|
|
|
show_gap = bool(head_rows) and bool(tail_rows)
|
|
return head_rows, tail_rows, show_gap
|
|
|
|
|
|
def _determine_preview_row_targets(num_rows: Optional[int]) -> Tuple[int, int]:
|
|
"""Compute how many head and tail rows to preview."""
|
|
max_rows = _DATASET_REPR_MAX_ROWS
|
|
if num_rows is None or num_rows <= max_rows:
|
|
head = num_rows if num_rows is not None else max_rows
|
|
return head, 0
|
|
|
|
head = min(_DATASET_REPR_HEAD_ROWS, max_rows)
|
|
tail = max_rows - head
|
|
return head, tail
|
|
|
|
|
|
def _format_rows_for_repr(
|
|
rows: List[Dict[str, Any]],
|
|
column_names: List[str],
|
|
) -> List[List[str]]:
|
|
"""Format row dicts into string cell rows for table rendering."""
|
|
formatted_rows: List[List[str]] = []
|
|
for row in rows:
|
|
formatted_row = []
|
|
for column in column_names:
|
|
value = row.get(column)
|
|
formatted_value = _format_value(value)
|
|
formatted_row.append(
|
|
_truncate_to_cell_width(formatted_value, _DATASET_REPR_MAX_COLUMN_WIDTH)
|
|
)
|
|
formatted_rows.append(formatted_row)
|
|
return formatted_rows
|
|
|
|
|
|
def _format_value(value: Any) -> str:
|
|
if isinstance(value, np.generic):
|
|
value = value.item()
|
|
return str(value).replace("\n", " ").replace("\r", " ")
|
|
|
|
|
|
def _truncate_to_cell_width(value: str, max_width: int) -> str:
|
|
"""Truncate a single cell to the configured max width."""
|
|
if max_width is None:
|
|
return value
|
|
if max_width <= 0:
|
|
return _DATASET_REPR_ELLIPSIS if value else ""
|
|
if len(value) <= max_width:
|
|
return value
|
|
if max_width == 1:
|
|
return _DATASET_REPR_ELLIPSIS
|
|
return value[: max_width - 1] + _DATASET_REPR_ELLIPSIS
|
|
|
|
|
|
def _compute_column_widths(
|
|
headers: List[str],
|
|
dtype_headers: List[str],
|
|
separator_row: List[str],
|
|
data_rows: List[List[str]],
|
|
) -> List[int]:
|
|
"""Compute per-column widths for table rendering."""
|
|
column_widths: List[int] = []
|
|
for idx in range(len(headers)):
|
|
widths = [
|
|
len(headers[idx]),
|
|
len(dtype_headers[idx]),
|
|
len(separator_row[idx]),
|
|
]
|
|
for row in data_rows:
|
|
widths.append(len(row[idx]))
|
|
column_widths.append(max(widths))
|
|
return column_widths
|
|
|
|
|
|
def _render_table_lines(
|
|
headers: List[str],
|
|
dtype_headers: List[str],
|
|
separator_row: List[str],
|
|
data_rows: List[List[str]],
|
|
column_widths: List[int],
|
|
) -> List[str]:
|
|
"""Render the full table (borders, headers, data) as lines."""
|
|
lines: List[str] = []
|
|
top = _render_border("╭", "┬", "╮", "─", column_widths)
|
|
header_row = _render_row(headers, column_widths)
|
|
separator_line = _render_row(separator_row, column_widths)
|
|
dtype_row = _render_row(dtype_headers, column_widths)
|
|
lines.extend([top, header_row, separator_line, dtype_row])
|
|
|
|
if data_rows:
|
|
middle = _render_border("╞", "╪", "╡", "═", column_widths)
|
|
lines.append(middle)
|
|
for row in data_rows:
|
|
lines.append(_render_row(row, column_widths))
|
|
|
|
bottom = _render_border("╰", "┴", "╯", "─", column_widths)
|
|
lines.append(bottom)
|
|
return lines
|
|
|
|
|
|
def _render_border(
|
|
left: str, middle: str, right: str, fill: str, column_widths: List[int]
|
|
) -> str:
|
|
"""Render a table border line given column widths."""
|
|
segments = [fill * (width + 2) for width in column_widths]
|
|
return f"{left}{middle.join(segments)}{right}"
|
|
|
|
|
|
def _render_row(values: List[str], column_widths: List[int]) -> str:
|
|
"""Render a single table row with padding."""
|
|
cells = []
|
|
for idx, value in enumerate(values):
|
|
padded = value.ljust(column_widths[idx])
|
|
cells.append(f" {padded} ")
|
|
return f"│{'┆'.join(cells)}│"
|
|
|
|
|
|
def _format_operator_dag(
|
|
op: "Operator",
|
|
curr_str: str = "",
|
|
depth: int = 0,
|
|
including_source: bool = True,
|
|
show_op_repr: bool = False,
|
|
) -> Tuple[str, int]:
|
|
"""Traverse (DFS) the Plan DAG and
|
|
return a string representation of the operators."""
|
|
if not including_source and isinstance(op, SourceOperator):
|
|
return curr_str, depth
|
|
|
|
curr_max_depth = depth
|
|
|
|
# For logical plan, only show the operator name like "Aggregate".
|
|
# But for physical plan, show the operator class name as well like
|
|
# "AllToAllOperator[Aggregate]".
|
|
op_str = repr(op) if show_op_repr else op.name
|
|
|
|
if depth == 0:
|
|
curr_str += f"{op_str}\n"
|
|
else:
|
|
trailing_space = " " * ((depth - 1) * 3)
|
|
curr_str += f"{trailing_space}+- {op_str}\n"
|
|
|
|
for input in op.input_dependencies:
|
|
curr_str, input_max_depth = _format_operator_dag(
|
|
input, curr_str, depth + 1, including_source, show_op_repr
|
|
)
|
|
curr_max_depth = max(curr_max_depth, input_max_depth)
|
|
return curr_str, curr_max_depth
|
|
|
|
|
|
def build_dataset_summary_repr(dataset: "Dataset") -> str:
|
|
"""Create a cosmetic string representation of a dataset.
|
|
|
|
This is used for Dataset.__repr__ when no tabular preview is available.
|
|
Must be very cheap — never forces execution.
|
|
"""
|
|
from ray.data.dataset import MaterializedDataset
|
|
|
|
dataset_cls = type(dataset)
|
|
logical_plan = dataset._logical_plan
|
|
dataset_name = dataset._dataset_name
|
|
|
|
plan_str = ""
|
|
plan_max_depth = 0
|
|
|
|
if not dataset._has_computed_output():
|
|
plan_str, plan_max_depth = _format_operator_dag(
|
|
logical_plan.dag, including_source=False
|
|
)
|
|
|
|
schema = dataset._base_schema(fetch_if_missing=False)
|
|
count = dataset._cache.get_num_rows(logical_plan.dag)
|
|
|
|
if schema is None or count is None:
|
|
has_n_ary_operator = False
|
|
dag = logical_plan.dag
|
|
|
|
while not isinstance(dag, SourceOperator):
|
|
if len(dag.input_dependencies) > 1:
|
|
has_n_ary_operator = True
|
|
break
|
|
|
|
dag = dag.input_dependencies[0]
|
|
|
|
# TODO(@bveeramani): Handle schemas for n-ary operators like `Union`.
|
|
if not has_n_ary_operator:
|
|
assert isinstance(dag, SourceOperator), dag
|
|
# We infer from logical plan's dag directly as we know that
|
|
# we don't have any cached values, so inferring is the only
|
|
# option left.
|
|
if schema is None:
|
|
schema = dag.infer_schema()
|
|
if count is None:
|
|
count = dag.infer_metadata().num_rows
|
|
|
|
if schema is None:
|
|
schema_str = "Unknown schema"
|
|
elif isinstance(schema, type):
|
|
schema_str = str(schema)
|
|
else:
|
|
schema_str = []
|
|
for n, t in zip(schema.names, schema.types):
|
|
if hasattr(t, "__name__"):
|
|
t = t.__name__
|
|
schema_str.append(f"{n}: {t}")
|
|
schema_str = ", ".join(schema_str)
|
|
schema_str = "{" + schema_str + "}"
|
|
|
|
if count is None:
|
|
count = "?"
|
|
|
|
num_blocks = None
|
|
if dataset_cls == MaterializedDataset:
|
|
num_blocks = logical_plan.initial_num_blocks()
|
|
assert num_blocks is not None
|
|
|
|
name_str = "name={}, ".format(dataset_name) if dataset_name is not None else ""
|
|
num_blocks_str = f"num_blocks={num_blocks}, " if num_blocks else ""
|
|
|
|
dataset_str = "{}({}{}num_rows={}, schema={})".format(
|
|
dataset_cls.__name__,
|
|
name_str,
|
|
num_blocks_str,
|
|
count,
|
|
schema_str,
|
|
)
|
|
|
|
# If the resulting string representation fits in one line, use it directly.
|
|
SCHEMA_LINE_CHAR_LIMIT = 80
|
|
MIN_FIELD_LENGTH = 10
|
|
INDENT_STR = " " * 3
|
|
trailing_space = INDENT_STR * plan_max_depth
|
|
|
|
if len(dataset_str) > SCHEMA_LINE_CHAR_LIMIT:
|
|
# If the resulting string representation exceeds the line char limit,
|
|
# first try breaking up each `Dataset` parameter into its own line
|
|
# and check if each line fits within the line limit. We check the
|
|
# `schema` param's length, since this is likely the longest string.
|
|
schema_str_on_new_line = f"{trailing_space}{INDENT_STR}schema={schema_str}"
|
|
if len(schema_str_on_new_line) > SCHEMA_LINE_CHAR_LIMIT:
|
|
# If the schema cannot fit on a single line, break up each field
|
|
# into its own line.
|
|
schema_str = []
|
|
for n, t in zip(schema.names, schema.types):
|
|
if hasattr(t, "__name__"):
|
|
t = t.__name__
|
|
col_str = f"{trailing_space}{INDENT_STR * 2}{n}: {t}"
|
|
# If the field line exceeds the char limit, abbreviate
|
|
# the field name to fit while maintaining the full type
|
|
if len(col_str) > SCHEMA_LINE_CHAR_LIMIT:
|
|
shortened_suffix = f"...: {str(t)}"
|
|
# Show at least 10 characters of the field name, even if
|
|
# we have already hit the line limit with the type.
|
|
chars_left_for_col_name = max(
|
|
SCHEMA_LINE_CHAR_LIMIT - len(shortened_suffix),
|
|
MIN_FIELD_LENGTH,
|
|
)
|
|
col_str = f"{col_str[:chars_left_for_col_name]}{shortened_suffix}"
|
|
schema_str.append(col_str)
|
|
schema_str = ",\n".join(schema_str)
|
|
schema_str = "{\n" + schema_str + f"\n{trailing_space}{INDENT_STR}" + "}"
|
|
name_str = (
|
|
f"\n{trailing_space}{INDENT_STR}name={dataset_name},"
|
|
if dataset_name is not None
|
|
else ""
|
|
)
|
|
num_blocks_str = (
|
|
f"\n{trailing_space}{INDENT_STR}num_blocks={num_blocks},"
|
|
if num_blocks
|
|
else ""
|
|
)
|
|
dataset_str = (
|
|
f"{dataset_cls.__name__}("
|
|
f"{name_str}"
|
|
f"{num_blocks_str}"
|
|
f"\n{trailing_space}{INDENT_STR}num_rows={count},"
|
|
f"\n{trailing_space}{INDENT_STR}schema={schema_str}"
|
|
f"\n{trailing_space})"
|
|
)
|
|
|
|
if plan_max_depth == 0:
|
|
plan_str += dataset_str
|
|
else:
|
|
plan_str += f"{INDENT_STR * (plan_max_depth - 1)}+- {dataset_str}"
|
|
return plan_str
|