Prefetch byte ranges using kvikio

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

@@ -6,8 +6,13 @@
 
 import concurrent.futures
 import contextlib
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+try:  # pragma: no cover; kvikio is optional
+    import kvikio
+except ImportError:
+    kvikio = None
 
 import pylibcudf as plc
 
@@ -46,6 +51,41 @@ class CachedParquetInfo:
     path: str
     size: int | None
     file_metadata: plc.io.parquet_metadata.FileMetaData
+    # Pre-created during footer prefetch; shared across all splits and scans of this file.
+    # HybridScanReader is not cached: it holds mutable per-read state so each worker
+    # creates its own from the shared metadata.
+    _hybrid_scan_metadata: list[plc.io.experimental.HybridScanMetadata] = field(
+        default_factory=list, compare=False, repr=False
+    )
+    _remote_handle: list[Any] = field(default_factory=list, compare=False, repr=False)
+
+    def hybrid_scan_reader(  # pragma: no cover; only called from thread pool workers where coverage.py does not trace
+        self,
+        options: plc.io.parquet.ParquetReaderOptions,
+    ) -> plc.io.experimental.HybridScanReader:
+        """Return a fresh HybridScanReader backed by shared pre-parsed file metadata."""
+        if not self._hybrid_scan_metadata:
+            self._hybrid_scan_metadata.append(
+                plc.io.experimental.HybridScanMetadata.from_parquet_metadata(
+                    self.file_metadata, options
+                )
+            )
+        return plc.io.experimental.HybridScanReader.from_metadata(
+            self._hybrid_scan_metadata[0]
+        )
+
+    def remote_handle(self) -> Any:  # pragma: no cover; requires kvikio
+        """Return the kvikio handle for this file."""
+        if not self._remote_handle:
+            if kvikio is None:
+                raise ImportError("kvikio is required for hybrid scan prefetching")
+            if plc.io.SourceInfo._is_remote_uri(self.path):
+                self._remote_handle.append(
+                    kvikio.RemoteFile.open(self.path, nbytes=self.size)
+                )
+            else:
+                self._remote_handle.append(kvikio.CuFile(self.path))
+        return self._remote_handle[0]
 
 
 @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths")
@@ -72,15 +112,10 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI
     # For now, we'll just use kvikio to explicitly get the size.
     sizes: list[int | None] = []
 
-    try:  # pragma: no cover; kvikio is optional
-        import kvikio
-    except ImportError:
-        kvikio = None
-
     for path in paths:
-        if (
-            paths and kvikio is not None and plc.io.SourceInfo._is_remote_uri(path)
-        ):  # pragma: no cover; kvikio is optional
+        if kvikio is not None and plc.io.SourceInfo._is_remote_uri(
+            path
+        ):  # pragma: no cover
             # We're OK to use `kvikio.RemoteFile.open` here. It does make an HTTP HEAD
             # request for S3/HTTP endpoints, but that's the entire reason we're running
             # this code. So long as it makes just *one* HTTP request, there's no advantage
@@ -99,10 +134,29 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI
         )
     )
 
-    return [
+    infos = [
         CachedParquetInfo(path, size, file_metadata)
         for path, size, file_metadata in zip(paths, sizes, metadata, strict=True)
     ]
+    for info in infos:
+        options = (
+            plc.io.parquet.ParquetReaderOptions.builder(
+                plc.io.SourceInfo([plc.io.types.FilepathSource(info.path, info.size)])
+            )
+            .decimal_width(plc.TypeId.DECIMAL128)
+            .build()
+        )
+        info._hybrid_scan_metadata.append(
+            plc.io.experimental.HybridScanMetadata.from_parquet_metadata(
+                info.file_metadata, options
+            )
+        )
+    if kvikio is not None:  # pragma: no cover; requires kvikio
+        # Open kvikio handles eagerly on the main thread before any prefetch workers
+        # start, so all splits sharing a file get the same handle without races.
+        for info in infos:
+            info.remote_handle()
+    return infos
 
 
 @nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir")