Prefetch parquet metadata for scan tasks

python/cudf_polars/cudf_polars/callback.py

@@ -24,7 +24,9 @@
 from rmm._cuda import gpu
 
 import cudf_polars.dsl.tracing
-from cudf_polars.dsl.ir import IRExecutionContext
+from cudf_polars.dsl.ir import (
+    IRExecutionContext,
+)
 from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN
 from cudf_polars.dsl.translate import Translator
 from cudf_polars.utils.config import (

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -20,6 +20,7 @@
 import itertools
 import json
 import random
+import reprlib
 import time
 import uuid
 from collections.abc import Sized
@@ -66,6 +67,7 @@
 )
 
 if TYPE_CHECKING:
+    import concurrent.futures
     from collections.abc import (
         Callable,
         Generator,
@@ -74,14 +76,14 @@
         Iterator,
         Sequence,
     )
-    from concurrent.futures import ThreadPoolExecutor
     from typing import Literal, Self
 
     from polars import polars  # type: ignore[attr-defined]
 
     from rmm.pylibrmm.stream import Stream
 
     from cudf_polars.containers.dataframe import NamedColumn
+    from cudf_polars.dsl.utils.io import CachedParquetInfo
     from cudf_polars.streaming.rank_aware_source import RankAwareSource
     from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice
     from cudf_polars.utils.config import ParquetOptions
@@ -138,7 +140,7 @@ class IRExecutionContext:
         Identifier for the query being executed.
     """
 
-    py_executor: ThreadPoolExecutor | None = field(default=None)
+    py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None)
     get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream)
     query_id: uuid.UUID = field(default_factory=uuid.uuid4)
 
@@ -570,6 +572,8 @@ def do_evaluate(
 def _parquet_physical_types(
     paths: list[str], columns: list[str] | None
 ) -> dict[str, plc.DataType]:
+    # TODO: Use prefetched metadata
+    # https://github.com/rapidsai/cudf/issues/22940
     metadata = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths))
     column_types = metadata.schema().column_types()
 
@@ -637,6 +641,7 @@ class Scan(IR):
     """Input from files."""
 
     __slots__ = (
+        "cached_parquet_info",
         "cloud_options",
         "include_file_paths",
         "n_rows",
@@ -663,7 +668,7 @@ class Scan(IR):
         "predicate",
         "parquet_options",
     )
-    _n_non_child_args = 11
+    _n_non_child_args = 12
     typ: str
     """What type of file are we reading? Parquet, CSV, etc..."""
     reader_options: dict[str, Any]
@@ -686,6 +691,8 @@ class Scan(IR):
     """Mask to apply to the read dataframe."""
     parquet_options: ParquetOptions
     """Parquet-specific options."""
+    cached_parquet_info: list[CachedParquetInfo] | None
+    """Cached parquet file metadata."""
 
     PARQUET_DEFAULT_CHUNK_SIZE: int = 0  # unlimited
     PARQUET_DEFAULT_PASS_LIMIT: int = 16 * 1024**3  # 16GiB
@@ -704,6 +711,7 @@ def __init__(
         include_file_paths: str | None,
         predicate: expr.NamedExpr | None,
         parquet_options: ParquetOptions,
+        cached_parquet_info: list[CachedParquetInfo] | None = None,
     ):
         self.schema = schema
         self.typ = typ
@@ -728,9 +736,14 @@ def __init__(
             include_file_paths,
             predicate,
             parquet_options,
+            cached_parquet_info,
         )
         self.children = ()
         self.parquet_options = parquet_options
+        self.cached_parquet_info = cached_parquet_info
+
+        Scan._validate_cached_parquet_info(self.paths, self.cached_parquet_info)
+
         if self.typ not in ("csv", "parquet", "ndjson"):  # pragma: no cover
             # This line is unhittable ATM since IPC/Anonymous scan raise
             # on the polars side
@@ -810,6 +823,21 @@ def __init__(
                 "Reading only parquet metadata to produce row index."
             )
 
+    @staticmethod
+    def _validate_cached_parquet_info(
+        paths: list[str],
+        cached_parquet_info: list[CachedParquetInfo] | None,
+    ) -> None:
+        if cached_parquet_info is not None and paths != [
+            info.path for info in cached_parquet_info
+        ]:
+            missing = reprlib.repr(
+                set(paths) - {info.path for info in cached_parquet_info}
+            )
+            raise AssertionError(
+                f"Paths do not match cached parquet info. Missing paths: {missing}"
+            )
+
     def get_hashable(self) -> Hashable:
         """
         Hashable representation of the node.
@@ -866,12 +894,34 @@ def add_file_paths(
     @staticmethod
     @nvtx_annotate_cudf_polars(message="Scan._get_parquet_row_count_from_metadata")
     def _get_parquet_row_count_from_metadata(
-        paths: list[str], skip_rows: int, n_rows: int
+        paths: list[str],
+        skip_rows: int,
+        n_rows: int,
+        parquet_options: ParquetOptions,
+        cached_parquet_info: list[CachedParquetInfo] | None,
     ) -> int:
         # Zero-width parquet files lose their row count when read through
         # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428
-        meta = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths))
-        num_rows = meta.num_rows() - skip_rows
+        if parquet_options.prefetch_file_metadata:
+            if cached_parquet_info is None:
+                raise AssertionError(
+                    "Cached parquet info is required when prefetching file metadata is enabled"
+                )
+
+            Scan._validate_cached_parquet_info(paths, cached_parquet_info)
+            parquet_metadatas = [
+                info.file_metadata for info in cached_parquet_info
+            ]  # pragma: no cover
+            num_rows = sum(
+                metadata.num_rows for metadata in parquet_metadatas
+            )  # pragma: no cover
+        else:
+            meta = plc.io.parquet_metadata.read_parquet_metadata(
+                plc.io.SourceInfo(paths)
+            )
+            num_rows = meta.num_rows()
+
+        num_rows -= skip_rows
         if n_rows != -1:
             num_rows = min(num_rows, n_rows)
         return max(num_rows, 0)
@@ -892,6 +942,7 @@ def do_evaluate(
         include_file_paths: str | None,
         predicate: expr.NamedExpr | None,
         parquet_options: ParquetOptions,
+        cached_parquet_info: list[CachedParquetInfo] | None,
         *,
         context: IRExecutionContext,
     ) -> DataFrame:
@@ -1007,6 +1058,24 @@ def read_csv_header(
                     df,
                 )
         elif typ == "parquet":
+            if parquet_options.prefetch_file_metadata:
+                if cached_parquet_info is None:
+                    raise AssertionError(
+                        "Cached parquet info is required when prefetching file metadata is enabled"
+                    )
+                Scan._validate_cached_parquet_info(paths, cached_parquet_info)
+                filepath_sources = []
+                parquet_metadatas = []
+                for info in cached_parquet_info:
+                    filepath_sources.append(
+                        plc.io.types.FilepathSource(info.path, info.size)
+                    )
+                    parquet_metadatas.append(info.file_metadata)
+                source_info = plc.io.SourceInfo(filepath_sources)
+            else:
+                parquet_metadatas = None
+                source_info = plc.io.SourceInfo(paths)
+
             filters = None
             if predicate is not None and row_index is None:
                 # Can't apply filters during read if we have a row index.
@@ -1016,9 +1085,7 @@ def read_csv_header(
                     ),
                     stream=stream,
                 )
-            builder = plc.io.parquet.ParquetReaderOptions.builder(
-                plc.io.SourceInfo(paths)
-            )
+            builder = plc.io.parquet.ParquetReaderOptions.builder(source_info)
             if filters is not None and parquet_options.use_jit_filter:
                 builder.use_jit_filter(use_jit_filter=True)
             parquet_reader_options = builder.decimal_width(
@@ -1038,6 +1105,7 @@ def read_csv_header(
                     parquet_reader_options,
                     chunk_read_limit=parquet_options.chunk_read_limit,
                     pass_read_limit=parquet_options.pass_read_limit,
+                    parquet_metadatas=parquet_metadatas,
                     stream=stream,
                 )
                 chunk = reader.read_chunk()
@@ -1053,7 +1121,9 @@ def read_csv_header(
                             [concatenated_columns[i], columns.pop()], stream=stream
                         )
                 num_rows = (
-                    cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows)
+                    cls._get_parquet_row_count_from_metadata(
+                        paths, skip_rows, n_rows, parquet_options, cached_parquet_info
+                    )
                     if not names
                     else None
                 )
@@ -1070,12 +1140,16 @@ def read_csv_header(
                     )
             else:
                 tbl_w_meta = plc.io.parquet.read_parquet(
-                    parquet_reader_options, stream=stream
+                    parquet_reader_options,
+                    parquet_metadatas=parquet_metadatas,
+                    stream=stream,
                 )
                 # TODO: consider nested column names?
                 col_names = tbl_w_meta.column_names(include_children=False)
                 num_rows = (
-                    cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows)
+                    cls._get_parquet_row_count_from_metadata(
+                        paths, skip_rows, n_rows, parquet_options, cached_parquet_info
+                    )
                     if not col_names
                     else None
                 )
@@ -1747,7 +1821,11 @@ def evaluate(
             stream = context.get_cuda_stream()
             scan = self.children[0]
             effective_rows = Scan._get_parquet_row_count_from_metadata(
-                scan.paths, scan.skip_rows, scan.n_rows
+                scan.paths,
+                scan.skip_rows,
+                scan.n_rows,
+                scan.parquet_options,
+                None,
             )
             dtype = DataType(pl.UInt32())
             col = Column(

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -466,6 +466,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I
             )
         ),
         parquet_options,
+        cached_parquet_info=None,
     )