chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,390 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
|
||||
|
||||
from ray._common.retry import call_with_retry
|
||||
from ray._private.utils import INT32_MAX
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import (
|
||||
reorder_columns_by_schema,
|
||||
)
|
||||
from ray.data._internal.execution.interfaces import TaskContext
|
||||
from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME
|
||||
from ray.data._internal.savemode import SaveMode
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.data.datasource.file_based_datasource import _resolve_kwargs
|
||||
from ray.data.datasource.file_datasink import _FileDatasink
|
||||
from ray.data.datasource.filename_provider import FilenameProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow
|
||||
|
||||
WRITE_FILE_MAX_ATTEMPTS = 10
|
||||
WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS = 32
|
||||
FILE_FORMAT = "parquet"
|
||||
|
||||
# These args are part of https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html#pyarrow.fs.FileSystem.open_output_stream
|
||||
# and are not supported by ParquetDatasink.
|
||||
UNSUPPORTED_OPEN_STREAM_ARGS = {"path", "buffer", "metadata"}
|
||||
|
||||
# https://arrow.apache.org/docs/python/generated/pyarrow.dataset.write_dataset.html
|
||||
ARROW_DEFAULT_MAX_ROWS_PER_GROUP = 1024 * 1024
|
||||
|
||||
DEFAULT_PARTITIONING_FLAVOR = "hive"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def choose_row_group_limits(
|
||||
row_group_size: Optional[int],
|
||||
min_rows_per_file: Optional[int],
|
||||
max_rows_per_file: Optional[int],
|
||||
) -> tuple[Optional[int], Optional[int], Optional[int]]:
|
||||
"""Configure row-group limits for Pyarrow's ``write_dataset`` API.
|
||||
|
||||
Configures the ``min_rows_per_group``, ``max_rows_per_group``, and
|
||||
``max_rows_per_file`` parameters based on Ray Data's configuration.
|
||||
|
||||
Args:
|
||||
row_group_size: The requested row-group size.
|
||||
min_rows_per_file: The minimum number of rows per file.
|
||||
max_rows_per_file: The maximum number of rows per file.
|
||||
|
||||
Returns:
|
||||
A tuple of (min_rows_per_group, max_rows_per_group, max_rows_per_file).
|
||||
"""
|
||||
|
||||
if (
|
||||
row_group_size is None
|
||||
and min_rows_per_file is None
|
||||
and max_rows_per_file is None
|
||||
):
|
||||
return None, None, None
|
||||
|
||||
elif row_group_size is None:
|
||||
# No explicit row group size provided. We are defaulting to
|
||||
# either the caller's min_rows_per_file or max_rows_per_file limits
|
||||
# or Arrow's defaults
|
||||
min_rows_per_group, max_rows_per_group, max_rows_per_file = (
|
||||
min_rows_per_file,
|
||||
max_rows_per_file,
|
||||
max_rows_per_file,
|
||||
)
|
||||
|
||||
# If min_rows_per_group is provided and max_rows_per_group is not,
|
||||
# and min_rows_per_group is greater than Arrow's default max_rows_per_group,
|
||||
# we set max_rows_per_group to min_rows_per_group to avoid creating too many row groups.
|
||||
if (
|
||||
min_rows_per_group is not None
|
||||
and max_rows_per_group is None
|
||||
and min_rows_per_group > ARROW_DEFAULT_MAX_ROWS_PER_GROUP
|
||||
):
|
||||
max_rows_per_group, max_rows_per_file = (
|
||||
min_rows_per_group,
|
||||
min_rows_per_group,
|
||||
)
|
||||
|
||||
return min_rows_per_group, max_rows_per_group, max_rows_per_file
|
||||
|
||||
elif row_group_size is not None and (
|
||||
min_rows_per_file is None or max_rows_per_file is None
|
||||
):
|
||||
return row_group_size, row_group_size, max_rows_per_file
|
||||
|
||||
else:
|
||||
# Clamp the requested `row_group_size` so that it is
|
||||
# * no smaller than `min_rows_per_file` (`lower`)
|
||||
# * no larger than `max_rows_per_file` (or Arrow's default cap) (`upper`)
|
||||
# This keeps each row-group within the per-file limits while staying
|
||||
# as close as possible to the requested size.
|
||||
clamped_group_size = max(
|
||||
min_rows_per_file, min(row_group_size, max_rows_per_file)
|
||||
)
|
||||
return clamped_group_size, clamped_group_size, max_rows_per_file
|
||||
|
||||
|
||||
def _widen_offset_overflowing_columns(
|
||||
tables: List["pyarrow.Table"], schema: "pyarrow.Schema"
|
||||
) -> "pyarrow.Schema":
|
||||
"""Promote `string`/`binary` columns to 64-bit-offset variants when needed.
|
||||
|
||||
Arrow addresses the data of `string` and `binary` columns with int32
|
||||
offsets, so a single contiguous array can hold at most `INT32_MAX` bytes.
|
||||
When the writer coalesces blocks into a large row group (e.g. because
|
||||
`min_rows_per_file` set `min_rows_per_group`), pyarrow must materialize each
|
||||
column of that row group as one contiguous array. If a `string`/`binary`
|
||||
column's combined size across the blocks in this write exceeds the int32
|
||||
limit, `write_dataset` fails with an offset-overflow error that surfaces as
|
||||
a column-length mismatch (the column is truncated to what fits).
|
||||
|
||||
Promoting such columns to `large_string`/`large_binary` (int64 offsets)
|
||||
removes the ceiling. The promotion is invisible on disk -- parquet stores
|
||||
both as `BYTE_ARRAY` -- so only the in-memory Arrow type changes.
|
||||
|
||||
Only top-level `string`/`binary` columns are promoted. Other int32-offset
|
||||
types are not handled: `string`/`binary` nested inside `list`/`struct`/`map`
|
||||
(whose inner offsets carry the same byte ceiling) and `list`/`map` columns
|
||||
themselves (which overflow on cumulative child-element count rather than
|
||||
bytes). These require recursive type rewriting and are far rarer in practice.
|
||||
|
||||
Args:
|
||||
tables: The blocks about to be written together in one task.
|
||||
schema: The unified output schema for those blocks.
|
||||
|
||||
Returns:
|
||||
``schema`` unchanged when no column overflows, otherwise a copy with the
|
||||
overflowing variable-width columns promoted to their ``large_*`` type.
|
||||
"""
|
||||
import pyarrow as pa
|
||||
|
||||
candidate_types = {}
|
||||
for field in schema:
|
||||
if pa.types.is_string(field.type):
|
||||
candidate_types[field.name] = pa.large_string()
|
||||
elif pa.types.is_binary(field.type):
|
||||
candidate_types[field.name] = pa.large_binary()
|
||||
if not candidate_types:
|
||||
return schema
|
||||
|
||||
# `.nbytes` is O(1) buffer metadata, so summing across blocks is cheap.
|
||||
overflowing: set = set()
|
||||
combined_nbytes: Dict[str, int] = defaultdict(int)
|
||||
for table in tables:
|
||||
for name in candidate_types:
|
||||
idx = table.schema.get_field_index(name)
|
||||
if idx != -1:
|
||||
combined_nbytes[name] += table.column(idx).nbytes
|
||||
if combined_nbytes[name] > INT32_MAX:
|
||||
overflowing.add(name)
|
||||
|
||||
if not overflowing:
|
||||
return schema
|
||||
|
||||
new_fields = [
|
||||
field.with_type(candidate_types[field.name])
|
||||
if field.name in overflowing
|
||||
else field
|
||||
for field in schema
|
||||
]
|
||||
return pa.schema(new_fields, metadata=schema.metadata)
|
||||
|
||||
|
||||
class ParquetDatasink(_FileDatasink):
|
||||
def __init__(
|
||||
self,
|
||||
path: str,
|
||||
*,
|
||||
partition_cols: Optional[List[str]] = None,
|
||||
arrow_parquet_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
|
||||
arrow_parquet_args: Optional[Dict[str, Any]] = None,
|
||||
min_rows_per_file: Optional[int] = None,
|
||||
max_rows_per_file: Optional[int] = None,
|
||||
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
||||
try_create_dir: bool = True,
|
||||
open_stream_args: Optional[Dict[str, Any]] = None,
|
||||
filename_provider: Optional[FilenameProvider] = None,
|
||||
dataset_uuid: Optional[str] = None,
|
||||
mode: SaveMode = SaveMode.APPEND,
|
||||
):
|
||||
if arrow_parquet_args_fn is None:
|
||||
arrow_parquet_args_fn = lambda: {} # noqa: E731
|
||||
|
||||
if arrow_parquet_args is None:
|
||||
arrow_parquet_args = {}
|
||||
|
||||
self.arrow_parquet_args_fn = arrow_parquet_args_fn
|
||||
self.arrow_parquet_args = arrow_parquet_args
|
||||
self.min_rows_per_file = min_rows_per_file
|
||||
self.max_rows_per_file = max_rows_per_file
|
||||
self.partition_cols = partition_cols
|
||||
|
||||
if self.min_rows_per_file is not None and self.max_rows_per_file is not None:
|
||||
if self.min_rows_per_file > self.max_rows_per_file:
|
||||
raise ValueError(
|
||||
"min_rows_per_file must be less than or equal to max_rows_per_file"
|
||||
)
|
||||
|
||||
if open_stream_args is not None:
|
||||
intersecting_keys = UNSUPPORTED_OPEN_STREAM_ARGS.intersection(
|
||||
set(open_stream_args.keys())
|
||||
)
|
||||
if intersecting_keys:
|
||||
logger.warning(
|
||||
"open_stream_args contains unsupported arguments: %s. These arguments "
|
||||
"are not supported by ParquetDatasink. They will be ignored.",
|
||||
intersecting_keys,
|
||||
)
|
||||
|
||||
if "compression" in open_stream_args:
|
||||
self.arrow_parquet_args["compression"] = open_stream_args["compression"]
|
||||
|
||||
if ("partitioning_flavor" in self.arrow_parquet_args) or (
|
||||
self.arrow_parquet_args_fn is not None
|
||||
and "partitioning_flavor" in self.arrow_parquet_args_fn()
|
||||
):
|
||||
if self.partition_cols is None:
|
||||
raise ValueError(
|
||||
"partition_cols must be provided when partitioning_flavor is set."
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
path,
|
||||
filesystem=filesystem,
|
||||
try_create_dir=try_create_dir,
|
||||
open_stream_args=open_stream_args,
|
||||
filename_provider=filename_provider,
|
||||
dataset_uuid=dataset_uuid,
|
||||
file_format=FILE_FORMAT,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
def write(
|
||||
self,
|
||||
blocks: Iterable[Block],
|
||||
ctx: TaskContext,
|
||||
) -> None:
|
||||
import pyarrow as pa
|
||||
|
||||
blocks = list(blocks)
|
||||
|
||||
if all(BlockAccessor.for_block(block).num_rows() == 0 for block in blocks):
|
||||
return
|
||||
|
||||
blocks = [
|
||||
block for block in blocks if BlockAccessor.for_block(block).num_rows() > 0
|
||||
]
|
||||
|
||||
filename = self.filename_provider.get_filename_for_task(
|
||||
ctx.kwargs[WRITE_UUID_KWARG_NAME], ctx.task_idx
|
||||
)
|
||||
write_kwargs = _resolve_kwargs(
|
||||
self.arrow_parquet_args_fn, **self.arrow_parquet_args
|
||||
)
|
||||
user_schema = write_kwargs.pop("schema", None)
|
||||
|
||||
# For partitioning_flavor, if it's not provided, the default is "hive"
|
||||
# Otherwise, it follows pyarrow's behavior: None for directory,
|
||||
# "hive" for hive, and "filename" for FilenamePartitioning.
|
||||
partitioning_flavor = write_kwargs.pop(
|
||||
"partitioning_flavor", DEFAULT_PARTITIONING_FLAVOR
|
||||
)
|
||||
|
||||
def write_blocks_to_path():
|
||||
tables = [BlockAccessor.for_block(block).to_arrow() for block in blocks]
|
||||
if user_schema is None:
|
||||
output_schema = pa.unify_schemas([table.schema for table in tables])
|
||||
# Coalescing many blocks into one row group can push a
|
||||
# `string`/`binary` column past Arrow's 2 GiB int32-offset
|
||||
# limit; promote such columns to their `large_*` variant so the
|
||||
# contiguous row-group array can address all of its bytes.
|
||||
output_schema = _widen_offset_overflowing_columns(tables, output_schema)
|
||||
else:
|
||||
output_schema = user_schema
|
||||
|
||||
self._write_parquet_files(
|
||||
tables,
|
||||
filename,
|
||||
output_schema,
|
||||
ctx.kwargs[WRITE_UUID_KWARG_NAME],
|
||||
write_kwargs,
|
||||
partitioning_flavor,
|
||||
)
|
||||
|
||||
logger.debug(f"Writing {filename} file to {self.path}.")
|
||||
|
||||
call_with_retry(
|
||||
write_blocks_to_path,
|
||||
description=f"write '{filename}' to '{self.path}'",
|
||||
match=self._data_context.retried_io_errors,
|
||||
max_attempts=WRITE_FILE_MAX_ATTEMPTS,
|
||||
max_backoff_s=WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS,
|
||||
)
|
||||
|
||||
def _get_basename_template(self, filename: str, write_uuid: str) -> str:
|
||||
# Check if write_uuid is present in filename, add if missing
|
||||
if write_uuid not in filename and self.mode == SaveMode.APPEND:
|
||||
raise ValueError(
|
||||
f"Write UUID '{write_uuid}' missing from filename template '{filename}'. This could result in files being overwritten."
|
||||
f"Modify your FileNameProvider implementation to include the `write_uuid` into the filename template or change your write mode to SaveMode.OVERWRITE. "
|
||||
)
|
||||
# Check if filename is already templatized
|
||||
if "{i}" in filename:
|
||||
# Filename is already templatized, but may need file extension
|
||||
if FILE_FORMAT not in filename:
|
||||
# Add file extension to templatized filename
|
||||
basename_template = f"{filename}.{FILE_FORMAT}"
|
||||
else:
|
||||
# Already has extension, use as-is
|
||||
basename_template = filename
|
||||
elif FILE_FORMAT not in filename:
|
||||
# No extension and not templatized, add extension and template
|
||||
basename_template = f"{filename}-{{i}}.{FILE_FORMAT}"
|
||||
else:
|
||||
# TODO(@goutamvenkat-anyscale): Add a warning if you pass in a custom
|
||||
# filename provider and it isn't templatized.
|
||||
# Use pathlib.Path to properly handle filenames with dots
|
||||
filename_path = Path(filename)
|
||||
stem = filename_path.stem # filename without extension
|
||||
assert "." not in stem, "Filename should not contain a dot"
|
||||
suffix = filename_path.suffix # extension including the dot
|
||||
basename_template = f"{stem}-{{i}}{suffix}"
|
||||
return basename_template
|
||||
|
||||
def _write_parquet_files(
|
||||
self,
|
||||
tables: List["pyarrow.Table"],
|
||||
filename: str,
|
||||
output_schema: "pyarrow.Schema",
|
||||
write_uuid: str,
|
||||
write_kwargs: Dict[str, Any],
|
||||
partitioning_flavor: Optional[str],
|
||||
) -> None:
|
||||
import pyarrow.dataset as ds
|
||||
|
||||
# Make every incoming batch conform to the final schema *before* writing.
|
||||
# `pa.unify_schemas` above fixed column order from the first block.
|
||||
for idx, table in enumerate(tables):
|
||||
if output_schema and not table.schema.equals(output_schema):
|
||||
table = reorder_columns_by_schema(table, output_schema)
|
||||
table = table.cast(output_schema)
|
||||
tables[idx] = table
|
||||
|
||||
row_group_size = write_kwargs.pop("row_group_size", None)
|
||||
|
||||
# We set this to "overwrite_or_ignore", to avoid the race condition seen in parallel writes when this is set to "error". The driver already handles the save mode check in on_write_start.
|
||||
existing_data_behavior = "overwrite_or_ignore"
|
||||
|
||||
(
|
||||
min_rows_per_group,
|
||||
max_rows_per_group,
|
||||
max_rows_per_file,
|
||||
) = choose_row_group_limits(
|
||||
row_group_size,
|
||||
min_rows_per_file=self.min_rows_per_file,
|
||||
max_rows_per_file=self.max_rows_per_file,
|
||||
)
|
||||
|
||||
basename_template = self._get_basename_template(filename, write_uuid)
|
||||
|
||||
# Note that the driver already handles the save mode logic, checking if the directory exists and raising an error if it does on SaveMode.ERROR
|
||||
ds.write_dataset(
|
||||
data=tables,
|
||||
base_dir=self.path,
|
||||
schema=output_schema,
|
||||
basename_template=basename_template,
|
||||
filesystem=self.filesystem,
|
||||
partitioning=self.partition_cols,
|
||||
format=FILE_FORMAT,
|
||||
existing_data_behavior=existing_data_behavior,
|
||||
partitioning_flavor=partitioning_flavor,
|
||||
use_threads=True,
|
||||
min_rows_per_group=min_rows_per_group,
|
||||
max_rows_per_group=max_rows_per_group,
|
||||
max_rows_per_file=max_rows_per_file,
|
||||
file_options=ds.ParquetFileFormat().make_write_options(**write_kwargs),
|
||||
create_dir=self.try_create_dir,
|
||||
)
|
||||
|
||||
@property
|
||||
def min_rows_per_write(self) -> Optional[int]:
|
||||
return self.min_rows_per_file
|
||||
Reference in New Issue
Block a user