[Data] [11/11] Enable footer-based Parquet reads by default

python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py

@@ -19,7 +19,7 @@
     get_type_hints,
 )
 
-from ray.data._internal.util import GiB, MiB, infer_compression
+from ray.data._internal.util import MiB, infer_compression
 from ray.util.annotations import DeveloperAPI
 
 
@@ -54,18 +54,6 @@ class LineDelimitedFileChunkMetadata(ChunkMetadata):
     chunk_byte_end_idx: int
 
 
-class ParquetFileChunkMetadata(ChunkMetadata):
-    """Metadata for Parquet file chunks.
-
-    For a parquet file, the chunks are based on the total size of the file, not on the
-    underlying row groups. We will split a file into potentially many chunks of the
-    target chunk size. This may correspond to 0, 1, or more row groups per chunk.
-    """
-
-    chunk_idx: int
-    total_num_chunks: int
-
-
 class ParquetRowGroupChunkMetadata(ChunkMetadata):
     """Metadata for a Parquet chunk described by explicit row-group indices.
 
@@ -174,64 +162,3 @@ def generate_chunk_metadatas(
                     ),
                     chunk_size,
                 )
-
-
-@DeveloperAPI
-class ParquetFileChunker(FileChunker):
-    """File chunker for Parquet files.
-
-    This chunker splits Parquet files into an estimated number of chunks. We do not
-    fetch the metadata for the file, so we may overestimate the number of chunks
-    compared to the actual number of underlying row groups. The partitioner creates
-    groupings based on these estimates, and the reader fetches the metadata and
-    ensures that all row groups are read / any overestimated row groups are ignored.
-    """
-
-    # Chosen so that we can effectively chunk files but will not result in OOMs if
-    # the compression ratio is high.
-    #
-    # If the compression ratio is high and this chunk size is large, we end up with
-    # larger chunks than we need and reading can OOM. Reducing the chunk size gives
-    # better memory performance by reading a smaller fraction of row groups at a time.
-    #
-    # We also want to keep this large enough such that we do not end up reading too
-    # much data if we underestimate the number of chunks. If row groups are larger
-    # than the chunk size and we place many of them in the same read task, the total
-    # amount of data read might be larger than expected. By increasing the chunk
-    # size we are less likely to put many such row groups in the same task.
-    _DEFAULT_TARGET_CHUNK_SIZE = 1 * GiB
-
-    def __init__(self, target_chunk_size: Optional[int] = None):
-        from ray.data.context import DataContext
-
-        ctx = DataContext.get_current()
-        if target_chunk_size is not None:
-            self._target_chunk_size = target_chunk_size
-        elif ctx.parquet_chunker_target_chunk_size is not None:
-            self._target_chunk_size = ctx.parquet_chunker_target_chunk_size
-        else:
-            self._target_chunk_size = self._DEFAULT_TARGET_CHUNK_SIZE
-
-    def generate_chunk_metadatas(
-        self, path: str, file_size: int
-    ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]:
-        if file_size <= self._target_chunk_size:
-            # Do not chunk if the file is smaller than the target chunk size; when
-            # we read the file, this prevents additional metadata fetching since we
-            # want to read the entire file.
-            yield None, file_size
-            return
-
-        num_chunks = math.ceil(file_size / self._target_chunk_size)
-        for i in range(num_chunks):
-            chunk_start = self._target_chunk_size * i
-            chunk_end = min(self._target_chunk_size * (i + 1), file_size)
-            chunk_size = chunk_end - chunk_start
-            yield (
-                create_chunk_metadata(
-                    ParquetFileChunkMetadata,
-                    chunk_idx=i,
-                    total_num_chunks=num_chunks,
-                ),
-                chunk_size,
-            )

python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py

@@ -1,124 +1,17 @@
 """Parquet chunk helpers for DataSourceV2.
 
-Two chunking strategies live here. The size-based one maps planner chunk
-metadata (``ParquetFileChunkMetadata``) to row-group ranges. The footer-based
-one maps ``ParquetRowGroupChunkMetadata`` -- the explicit surviving row groups a
-bin assigns to a file -- straight to PyArrow ``ParquetFileFragment`` subsets.
-Both produce ``(fragment, file_row_offset)`` pairs for the reader.
+Maps ``ParquetRowGroupChunkMetadata`` (the explicit surviving row groups a bin
+assigns to a file) to PyArrow ``ParquetFileFragment`` subsets for reading.
 """
-from typing import Callable, Iterable, List, Optional, Tuple, TypeVar
+from typing import Callable, Iterable, List, Tuple, TypeVar
 
 import pyarrow.dataset as pds
 
 from ray._common.retry import call_with_retry
-from ray.data._internal.datasource_v2.chunkers.file_chunker import (
-    ParquetFileChunkMetadata,
-)
 
 R = TypeVar("R")
 
 
-def _calculate_row_group_range(
-    chunk_idx: int, total_num_chunks: int, total_row_groups: int
-) -> Optional[Tuple[int, int]]:
-    """Compute the half-open row-group range for a given chunk.
-
-    Distributes row groups as evenly as possible across chunks. If row groups
-    don't divide evenly, earlier chunks get the extra row groups.
-
-    Example:
-        - 10 row groups, 3 chunks -> [0:4), [4:7), [7:10)
-        - 11 row groups, 3 chunks -> [0:4), [4:8), [8:11)
-
-    Args:
-        chunk_idx: Index of the current chunk (0-based).
-        total_num_chunks: Total number of chunks.
-        total_row_groups: Total number of row groups to distribute.
-
-    Returns:
-        Tuple ``(start_row_group, end_row_group)`` where ``end`` is exclusive,
-        or ``None`` if ``chunk_idx`` falls beyond the actual number of row
-        groups (i.e. the planner over-estimated the chunk count).
-    """
-    assert (
-        total_row_groups >= 0
-    ), f"total_row_groups must be non-negative, got {total_row_groups}"
-    assert (
-        total_num_chunks > 0
-    ), f"total_num_chunks must be positive, got {total_num_chunks}"
-    assert (
-        chunk_idx < total_num_chunks
-    ), f"chunk_idx must be less than total_num_chunks, got {chunk_idx} and {total_num_chunks}"
-    assert chunk_idx >= 0, f"chunk_idx must be non-negative, got {chunk_idx}"
-
-    # Handle the case where ``chunk_idx`` exceeds the actual number of chunks
-    # needed. This happens when the planner overestimated the number of chunks
-    # (the chunker doesn't fetch metadata).
-    if chunk_idx >= total_row_groups:
-        return None
-
-    base_row_groups_per_chunk = total_row_groups // total_num_chunks
-    remainder = total_row_groups % total_num_chunks
-
-    # Chunks 0 through (remainder-1) get one extra row group.
-    if chunk_idx < remainder:
-        row_groups_in_this_chunk = base_row_groups_per_chunk + 1
-        start = chunk_idx * row_groups_in_this_chunk
-    else:
-        row_groups_in_this_chunk = base_row_groups_per_chunk
-        start = (
-            remainder * (base_row_groups_per_chunk + 1)
-            + (chunk_idx - remainder) * base_row_groups_per_chunk
-        )
-
-    end = start + row_groups_in_this_chunk
-
-    assert (
-        0 <= start <= end <= total_row_groups
-    ), f"Invalid range [{start}, {end}) for {total_row_groups} row groups"
-
-    return start, end
-
-
-def _fragments_from_chunk_metadata(
-    fragment: pds.ParquetFileFragment,
-    chunk_metadata: ParquetFileChunkMetadata,
-) -> List[Tuple[pds.ParquetFileFragment, int]]:
-    """Slice ``fragment`` into per-row-group sub-fragments per chunk metadata.
-
-    Returns one ``(ParquetFileFragment, file_row_offset)`` pair per row group
-    covered by the chunk, where ``file_row_offset`` is the sum of ``num_rows``
-    across all row groups that precede the sub-fragment in the underlying
-    file. Callers seed per-fragment hashing offsets with this value so
-    sub-fragments of the same file don't collide on ``(path, 0, n)``.
-
-    Returns an empty list when the chunk index falls beyond the file's actual
-    row-group count (the planner over-estimated; we silently drop the slice).
-    """
-    chunk_idx = chunk_metadata["chunk_idx"]
-    total_num_chunks = chunk_metadata["total_num_chunks"]
-    metadata = fragment.metadata
-    total_row_groups = metadata.num_row_groups
-
-    row_group_range = _calculate_row_group_range(
-        chunk_idx, total_num_chunks, total_row_groups
-    )
-
-    if row_group_range is None:
-        return []
-
-    start, end = row_group_range
-
-    file_row_offset = sum(metadata.row_group(i).num_rows for i in range(start))
-    sub_fragments: List[Tuple[pds.ParquetFileFragment, int]] = []
-    for row_group_index in range(start, end):
-        sub_fragments.append(
-            (fragment.subset(row_group_ids=[row_group_index]), file_row_offset)
-        )
-        file_row_offset += metadata.row_group(row_group_index).num_rows
-    return sub_fragments
-
-
 def _with_io_retry(f: Callable[[], R], description: str) -> R:
     """Run ``f``, retrying the transient IO errors configured on the context.
 

python/ray/data/_internal/datasource_v2/listing/file_indexer.py

@@ -460,9 +460,8 @@ def _process_file_infos_to_manifests(
             path, file_size = file_info.path, file_info.size
 
             # Drive the chunker once per file; emit one manifest row per chunk.
-            # ``chunk_metadata`` is ``None`` for whole-file chunks (default
-            # ``WholeFileChunker`` behavior and ``ParquetFileChunker`` for files
-            # smaller than the target chunk size).
+            # ``chunk_metadata`` is ``None`` for whole-file chunks (the default
+            # ``WholeFileChunker`` behavior).
             for (
                 chunk_metadata,
                 chunk_size,

python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py

@@ -21,18 +21,11 @@
     ParquetDatasource,
     check_for_legacy_tensor_type,
 )
-from ray.data._internal.datasource_v2.chunkers.file_chunker import (
-    FileChunker,
-    ParquetFileChunker,
-)
 from ray.data._internal.datasource_v2.datasource_v2 import (
     DatasourceCategory,
     DataSourceV2,
 )
-from ray.data._internal.datasource_v2.listing.file_indexer import (
-    FileIndexer,
-    NonSamplingFileIndexer,
-)
+from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer
 from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
 from ray.data._internal.datasource_v2.readers.file_reader import (
     INCLUDE_PATHS_COLUMN_NAME,
@@ -86,7 +79,6 @@ def __init__(
         arrow_parquet_args: Optional[dict] = None,
         schema: Optional[pa.Schema] = None,
         parquet_format_kwargs: Optional[dict] = None,
-        file_chunker: Optional[FileChunker] = None,
     ):
         super().__init__(name="ParquetV2", category=DatasourceCategory.FILE_BASED)
         # Capture the ``local://`` check against the *original* paths;
@@ -140,14 +132,6 @@ def __init__(
         # footers, and the scanner pins it on the pyarrow dataset so files
         # are cast to these types at scan time.
         self._user_schema = schema
-        # Chunker that splits each listed Parquet file into one or more
-        # row-group-aligned read units. Defaults to ``ParquetFileChunker``
-        # (1 GiB target chunk size, or whatever ``DataContext`` configures).
-        # Callers can inject an alternative for tests or shuffle-aware
-        # planning code that wants whole-file reads.
-        self._file_chunker: FileChunker = (
-            file_chunker if file_chunker is not None else ParquetFileChunker()
-        )
 
     @property
     def paths(self) -> List[str]:
@@ -182,56 +166,46 @@ def shuffle(self) -> Optional[Union[Literal["files"], "FileShuffleConfig"]]:
         return self._shuffle
 
     def _get_file_indexer(self) -> FileIndexer:
-        # Opt-in footer-based indexing: reads each file's footer on a
-        # ``FooterReader`` actor pool so listing rows carry per-row-group stats
-        # (and predicate / limit / projection push-down reach listing). Grouping
-        # those rows into read units is ``get_file_partitioner``'s job. Off by
-        # default while it is validated against release benchmarks; the flag and
-        # the blind chunker below both go away once it is the only path.
-        if env_bool("RAY_DATA_PARQUET_ENABLE_FOOTER_INDEXER", False):
-            from ray.data._internal.datasource_v2.listing.footer_file_indexer import (
-                FooterFileIndexer,
-            )
-
-            return FooterFileIndexer(
-                ignore_missing_paths=self._ignore_missing_paths,
-                skip_paths=self._skip_paths,
-                coalesce_bytes=env_integer("RAY_DATA_PARQUET_FOOTER_COALESCE_BYTES", 0),
-            )
+        # Parquet V2 reads always use the footer-based indexer: it reads each
+        # file's footer on a ``FooterReader`` actor pool so listing rows carry
+        # per-row-group stats, and predicate / limit / projection push-down
+        # reach listing. Grouping those rows into read units is
+        # ``get_file_partitioner``'s job.
+        from ray.data._internal.datasource_v2.listing.footer_file_indexer import (
+            FooterFileIndexer,
+        )
 
-        return NonSamplingFileIndexer(
+        return FooterFileIndexer(
             ignore_missing_paths=self._ignore_missing_paths,
             skip_paths=self._skip_paths,
-            file_chunker=self._file_chunker,
+            coalesce_bytes=env_integer("RAY_DATA_PARQUET_FOOTER_COALESCE_BYTES", 0),
         )
 
     def get_file_partitioner(self, **kwargs):
-        # With the footer indexer on, listing rows carry per-row-group stats, so
-        # bin-pack them into read units instead of size-estimating whole files.
-        if env_bool("RAY_DATA_PARQUET_ENABLE_FOOTER_INDEXER", False):
-            from ray.data._internal.datasource_v2.partitioners.online_bin_packer import (  # noqa: E501
-                OnlineBinPacker,
-            )
-            from ray.data._internal.util import MiB
+        # Listing rows carry per-row-group stats, so bin-pack them into read
+        # units instead of size-estimating whole files.
+        from ray.data._internal.datasource_v2.partitioners.online_bin_packer import (
+            OnlineBinPacker,
+        )
+        from ray.data._internal.util import MiB
 
-            max_bin_bytes = env_integer("RAY_DATA_PARQUET_BIN_PACKING_BYTES", 128 * MiB)
-            max_shared_open_bins = env_integer(
-                "RAY_DATA_PARQUET_BIN_PACKING_MAX_SHARED_OPEN_BINS", 16
-            )
-            split_coalesced = env_bool("RAY_DATA_PARQUET_FOOTER_SPLIT_COALESCED", False)
-            logger.debug(
-                "OnlineBinPacker(max_bin_bytes=%d, max_shared_open_bins=%d, "
-                "split_coalesced=%s)",
-                max_bin_bytes,
-                max_shared_open_bins,
-                split_coalesced,
-            )
-            return OnlineBinPacker(
-                max_bin_bytes=max_bin_bytes,
-                max_shared_open_bins=max_shared_open_bins,
-                split_coalesced=split_coalesced,
-            )
-        return super().get_file_partitioner(**kwargs)
+        max_bin_bytes = env_integer("RAY_DATA_PARQUET_BIN_PACKING_BYTES", 128 * MiB)
+        max_shared_open_bins = env_integer(
+            "RAY_DATA_PARQUET_BIN_PACKING_MAX_SHARED_OPEN_BINS", 16
+        )
+        split_coalesced = env_bool("RAY_DATA_PARQUET_FOOTER_SPLIT_COALESCED", False)
+        logger.debug(
+            "OnlineBinPacker(max_bin_bytes=%d, max_shared_open_bins=%d, "
+            "split_coalesced=%s)",
+            max_bin_bytes,
+            max_shared_open_bins,
+            split_coalesced,
+        )
+        return OnlineBinPacker(
+            max_bin_bytes=max_bin_bytes,
+            max_shared_open_bins=max_shared_open_bins,
+            split_coalesced=split_coalesced,
+        )
 
     def get_size_estimator(self) -> ParquetInMemorySizeEstimator:
         return ParquetInMemorySizeEstimator()