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