[PERF]: Improve the performance of plc.io.parquet.read_parquet with prefetched parquet file metadata.

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -902,12 +902,7 @@ def _get_parquet_row_count_from_metadata(
     ) -> int:
         # Zero-width parquet files lose their row count when read through
         # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428
-        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"
-                )
-
+        if cached_parquet_info is not None:
             Scan._validate_cached_parquet_info(paths, cached_parquet_info)
             parquet_metadatas = [
                 info.file_metadata for info in cached_parquet_info
@@ -1071,11 +1066,7 @@ 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"
-                    )
+            if cached_parquet_info is not None:
                 Scan._validate_cached_parquet_info(paths, cached_parquet_info)
                 filepath_sources = []
                 parquet_metadatas = []

python/cudf_polars/cudf_polars/dsl/utils/io.py

@@ -110,6 +110,8 @@ def prefetch_parquet_file_metadata_for_ir(
     root: IR,
     py_executor: concurrent.futures.Executor | None,
     stats: StatsCollector | None = None,
+    *,
+    remote_only: bool = False,
 ) -> dict[str, CachedParquetInfo]:
     """
     Prefetch parquet metadata for all parquet scans in an IR graph.
@@ -125,6 +127,9 @@ def prefetch_parquet_file_metadata_for_ir(
         prefetched during statistics collection, when the number of files
         sampled equals the total number of files. Providing ``stats`` here will
         skip rereading metadata for those files.
+    remote_only
+        If ``True``, only prefetch metadata for remote URIs (e.g. ``s3://``),
+        skipping local paths.
 
     Returns
     -------
@@ -135,7 +140,7 @@ def prefetch_parquet_file_metadata_for_ir(
     all_paths: set[str] = set()
 
     for node in traversal([root]):
-        if isinstance(node, StreamingScan):
+        if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet":
             for scan in node.scans:
                 for path in scan.paths:
                     all_paths.add(path)
@@ -155,6 +160,10 @@ def prefetch_parquet_file_metadata_for_ir(
                     cached_parquet_info[info.path] = info
 
     missing_paths = all_paths - set(cached_parquet_info.keys())
+    if remote_only:
+        missing_paths = {
+            p for p in missing_paths if plc.io.SourceInfo._is_remote_uri(p)
+        }
     cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None]
 
     if py_executor is None:
@@ -194,8 +203,10 @@ def attach_cached_parquet_metadata(
         Mapping from file paths to cached parquet metadata.
     """
     for node in traversal([root]):
-        if isinstance(node, StreamingScan):
+        if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet":
             for scan in node.scans:
+                if not all(path in cached_parquet_info_map for path in scan.paths):
+                    continue
                 cached = [cached_parquet_info_map[path] for path in scan.paths]
                 Scan._validate_cached_parquet_info(scan.paths, cached)
                 scan.cached_parquet_info = cached

python/cudf_polars/cudf_polars/engine/core.py

@@ -42,7 +42,7 @@
 from cudf_polars.streaming.parallel import lower_ir_graph_with_node_map
 from cudf_polars.streaming.statistics import collect_statistics
 from cudf_polars.streaming.utils import _concat
-from cudf_polars.utils.config import get_total_device_memory
+from cudf_polars.utils.config import Unspecified, get_total_device_memory
 
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping
@@ -778,11 +778,13 @@ def evaluate_on_rank(
         py_executor, get_cuda_stream=ctx.br().stream_pool.get_stream, query_id=query_id
     )
 
-    if config_options.parquet_options.prefetch_file_metadata:
+    prefetch_file_metadata = config_options.parquet_options.prefetch_file_metadata
+    if prefetch_file_metadata is not False:
         cached_parquet_info_map = prefetch_parquet_file_metadata_for_ir(
             ir,
             ir_context.py_executor,
             stats=stats,
+            remote_only=isinstance(prefetch_file_metadata, Unspecified),
         )
         attach_cached_parquet_metadata(ir, cached_parquet_info_map)
 

python/cudf_polars/cudf_polars/engine/options.py

@@ -18,7 +18,7 @@
 from cudf_polars.engine.hardware_binding import (
     HardwareBindingPolicy,
 )
-from cudf_polars.utils.config import MemoryResourceConfig
+from cudf_polars.utils.config import UNSPECIFIED, MemoryResourceConfig, Unspecified
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -37,38 +37,6 @@
 ]
 
 
-class Unspecified:
-    """
-    Sentinel value meaning "fall back to environment variable, then built-in default".
-
-    The singleton instance :data:`UNSPECIFIED` is used as the default for every
-    :class:`StreamingOptions` field.  When a field is still ``UNSPECIFIED`` after
-    construction (i.e. neither an explicit value nor an environment variable was provided),
-    the underlying library applies its own built-in default.
-    """
-
-    _instance: Unspecified | None = None
-
-    def __new__(cls) -> Unspecified:
-        """Return the singleton instance."""
-        if cls._instance is None:
-            cls._instance = super().__new__(cls)
-        return cls._instance
-
-    def __repr__(self) -> str:
-        """Return ``"UNSPECIFIED"``."""
-        return "UNSPECIFIED"
-
-
-UNSPECIFIED = Unspecified()
-"""Singleton sentinel for all :class:`StreamingOptions` fields.
-
-A field set to ``UNSPECIFIED`` after construction means no explicit value and no
-matching environment variable was found; the underlying library will apply its own
-built-in default.
-"""
-
-
 def _opt(
     category: str,
     env_var: str | None = None,

python/cudf_polars/cudf_polars/utils/config.py

@@ -49,6 +49,7 @@
 
 
 __all__ = [
+    "UNSPECIFIED",
     "Cluster",
     "ConfigOptions",
     "DaskContext",
@@ -60,9 +61,45 @@
     "SPMDContext",
     "StreamingExecutor",
     "StreamingFallbackMode",
+    "Unspecified",
 ]
 
 
+class Unspecified:
+    """
+    Sentinel value meaning "no value was explicitly provided".
+
+    The singleton instance :data:`UNSPECIFIED` is used as the default for every
+    :class:`StreamingOptions` field, as well as for
+    :attr:`ParquetOptions.prefetch_file_metadata`. When a field is still
+    ``UNSPECIFIED`` after construction (i.e. neither an explicit value nor a
+    matching environment variable was provided), the consuming component decides
+    on the semantics.
+    """
+
+    _instance: Unspecified | None = None
+
+    def __new__(cls) -> Unspecified:
+        """Return the singleton instance."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __repr__(self) -> str:
+        """Return ``"UNSPECIFIED"``."""
+        return "UNSPECIFIED"
+
+
+UNSPECIFIED = Unspecified()
+"""Singleton sentinel for all :class:`StreamingOptions` fields, as well as for
+:attr:`ParquetOptions.prefetch_file_metadata`.
+
+A field set to ``UNSPECIFIED`` after construction means no explicit value and no
+matching environment variable was found; the consuming component decides on the
+semantics.
+"""
+
+
 def _env_get_int(name: str, default: int) -> int:
     try:
         return int(os.getenv(name, default))
@@ -221,7 +258,10 @@ class ParquetOptions:
         will also be skipped if ``max_footer_samples`` is 0.
     prefetch_file_metadata
         Whether to prefetch parquet file metadata and pass it through
-        `parquet_metadatas` to avoid rereading file footers.
+        `parquet_metadatas` to avoid rereading file footers. Not supported
+        by the in-memory executor, where it defaults to disabled. For the
+        streaming executor, it defaults to being enabled for remote URIs
+        (e.g. ``s3://``) only; pass ``True`` to also prefetch local files.
     use_jit_filter
         Whether to use JIT compilation for post-read filtering in Parquet scans.
         When enabled, filter predicates are JIT-compiled to CUDA kernels for
@@ -261,11 +301,11 @@ class ParquetOptions:
             f"{_env_prefix}__MAX_ROW_GROUP_SAMPLES", int, default=1
         )
     )
-    prefetch_file_metadata: bool = dataclasses.field(
+    prefetch_file_metadata: bool | Unspecified = dataclasses.field(
         default_factory=_make_default_factory(
             f"{_env_prefix}__PREFETCH_FILE_METADATA",
             _bool_converter,
-            default=False,
+            default=UNSPECIFIED,
         )
     )
     use_jit_filter: bool = dataclasses.field(
@@ -289,8 +329,8 @@ def __post_init__(self) -> None:  # noqa: D105
             raise TypeError("max_footer_samples must be an int")
         if not isinstance(self.max_row_group_samples, int):
             raise TypeError("max_row_group_samples must be an int")
-        if not isinstance(self.prefetch_file_metadata, bool):
-            raise TypeError("prefetch_file_metadata must be a bool")
+        if not isinstance(self.prefetch_file_metadata, (bool, Unspecified)):
+            raise TypeError("prefetch_file_metadata must be a bool when specified")
         if not isinstance(self.use_jit_filter, bool):
             raise TypeError("use_jit_filter must be a bool")
 
@@ -960,9 +1000,31 @@ def from_polars_engine(
         if user_parquet_options is None:
             user_parquet_options = {}
 
+        # Engine-dependent default: only prefetch for the streaming executor.
+        # Skipped if the user or the environment has already set a value.
+        prefetch_default = UNSPECIFIED if user_executor == "streaming" else False
+        prefetch_env_set = (
+            os.environ.get(f"{ParquetOptions._env_prefix}__PREFETCH_FILE_METADATA")
+            is not None
+        )
+
         if isinstance(user_parquet_options, dict):
+            user_parquet_options = dict(user_parquet_options)
+            if (
+                "prefetch_file_metadata" not in user_parquet_options
+                and not prefetch_env_set
+            ):
+                user_parquet_options["prefetch_file_metadata"] = prefetch_default
             parquet_options = ParquetOptions(**user_parquet_options)
         else:
+            if (
+                isinstance(user_parquet_options.prefetch_file_metadata, Unspecified)
+                and not prefetch_env_set
+            ):
+                user_parquet_options = dataclasses.replace(
+                    user_parquet_options,
+                    prefetch_file_metadata=prefetch_default,
+                )
             parquet_options = user_parquet_options
         # This is set in polars, and so can't be overridden by the environment
         user_raise_on_fail = engine.config.get("raise_on_fail", False)
@@ -991,7 +1053,7 @@ def from_polars_engine(
         match user_executor:
             case "in-memory":
                 executor = InMemoryExecutor(**user_executor_options)
-                if parquet_options.prefetch_file_metadata:
+                if parquet_options.prefetch_file_metadata is True:
                     raise NotImplementedError(
                         "Prefetching is not supported for the in-memory executor."
                     )

python/cudf_polars/tests/streaming/test_scan.py

@@ -126,6 +126,27 @@ def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None:
     assert result == {}
 
 
+def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None:
+    make_partitioned_source(df, tmp_path, "parquet", n_files=1)
+    local_path = str(next(tmp_path.glob("*.parquet")))
+
+    scan = _make_parquet_scan([local_path])
+    fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, [])
+    streaming_scan = StreamingScan([fused], scan, "fused")
+
+    # Local paths are skipped entirely when remote_only=True.
+    result = prefetch_parquet_file_metadata_for_ir(
+        streaming_scan, py_executor=None, stats=None, remote_only=True
+    )
+    assert result == {}
+
+    # The same local path is prefetched when remote_only=False (the default).
+    result = prefetch_parquet_file_metadata_for_ir(
+        streaming_scan, py_executor=None, stats=None
+    )
+    assert set(result) == {local_path}
+
+
 def test_prefetch_file_metadata_select_fast_count(
     df: pl.DataFrame,
     streaming_engine_factory: Callable[..., StreamingEngine],
@@ -349,33 +370,13 @@ def test_streaming_scan_raises() -> None:
         StreamingScan.do_evaluate([fused], scan, context=ctx)
 
 
-def test_scan_missing_prefetch_metadata_raises() -> None:
+def test_scan_path_mismatch_raises() -> None:
     # This isn't reachable by polars' public API, so we test it directly.
     scan = _make_parquet_scan(
         ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True)
     )
     ctx = IRExecutionContext()
 
-    with pytest.raises(
-        AssertionError,
-        match=r"Cached parquet info is required",
-    ):
-        Scan.do_evaluate(
-            scan.schema,
-            scan.typ,
-            scan.reader_options,
-            scan.paths,
-            scan.with_columns,
-            scan.skip_rows,
-            scan.n_rows,
-            scan.row_index,
-            scan.include_file_paths,
-            scan.predicate,
-            scan.parquet_options,
-            None,
-            context=ctx,
-        )
-
     with pytest.raises(
         AssertionError,
         match=r"Paths do not match cached parquet info",