Implement cudf-polars chunked parquet reading

python/cudf_polars/cudf_polars/callback.py

@@ -129,6 +129,7 @@ def set_device(device: int | None) -> Generator[int, None, None]:
 
 def _callback(
     ir: IR,
+    config: GPUEngine,
     with_columns: list[str] | None,
     pyarrow_predicate: str | None,
     n_rows: int | None,
@@ -145,7 +146,7 @@ def _callback(
         set_device(device),
         set_memory_resource(memory_resource),
     ):
-        return ir.evaluate(cache={}).to_polars()
+        return ir.evaluate(cache={}, config=config).to_polars()
 
 
 def execute_with_cudf(
@@ -174,7 +175,7 @@ def execute_with_cudf(
     device = config.device
     memory_resource = config.memory_resource
     raise_on_fail = config.config.get("raise_on_fail", False)
-    if unsupported := (config.config.keys() - {"raise_on_fail"}):
+    if unsupported := (config.config.keys() - {"raise_on_fail", "parquet_options"}):
         raise ValueError(
             f"Engine configuration contains unsupported settings {unsupported}"
         )
@@ -184,6 +185,7 @@ def execute_with_cudf(
                 partial(
                     _callback,
                     translate_ir(nt),
+                    config,
                     device=device,
                     memory_resource=memory_resource,
                 )

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -34,8 +34,12 @@
     from collections.abc import Callable, Hashable, MutableMapping, Sequence
     from typing import Literal
 
+    from polars import GPUEngine
+
     from cudf_polars.typing import Schema
 
+PARQUET_DEFAULT_CHUNK_SIZE = 0
+PARQUET_DEFAULT_PASS_LIMIT = 17179869184  # 16GiB
 
 __all__ = [
     "IR",
@@ -144,7 +148,9 @@ def get_hashable(self) -> Hashable:
         schema_hash = tuple(self.schema.items())
         return (type(self), schema_hash, args)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """
         Evaluate the node and return a dataframe.
 
@@ -153,6 +159,8 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         cache
             Mapping from cached node ids to constructed DataFrames.
             Used to implement evaluation of the `Cache` node.
+        config
+            GPU engine configuration.
 
         Returns
         -------
@@ -339,7 +347,9 @@ def get_hashable(self) -> Hashable:
             self.predicate,
         )
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         with_columns = self.with_columns
         row_index = self.row_index
@@ -418,17 +428,47 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
                 colnames[0],
             )
         elif self.typ == "parquet":
-            tbl_w_meta = plc.io.parquet.read_parquet(
-                plc.io.SourceInfo(self.paths),
-                columns=with_columns,
-                nrows=n_rows,
-                skip_rows=self.skip_rows,
-            )
-            df = DataFrame.from_table(
-                tbl_w_meta.tbl,
-                # TODO: consider nested column names?
-                tbl_w_meta.column_names(include_children=False),
-            )
+            parquet_options = config.config.get("parquet_options", {})
+            if parquet_options.get("chunked", False):
+                reader = plc.io.parquet.ChunkedParquetReader(
+                    plc.io.SourceInfo(self.paths),
+                    columns=with_columns,
+                    nrows=n_rows,
+                    skip_rows=self.skip_rows,
+                    chunk_read_limit=parquet_options.get(
+                        "chunk_read_limit", PARQUET_DEFAULT_CHUNK_SIZE
+                    ),
+                    pass_read_limit=parquet_options.get(
+                        "pass_read_limit", PARQUET_DEFAULT_PASS_LIMIT
+                    ),
+                )
+
+                chunks = []  # type: ignore
+                chk = reader.read_chunk()
+                tbl = chk.tbl
+                chunks.append(tbl)
+                names = chk.column_names()
+                while reader.has_next():
+                    chunks.append(reader.read_chunk().tbl)
+
+                chunks = plc.concatenate.concatenate(chunks)
+                df = DataFrame.from_table(
+                    chunks,
+                    names=names,
+                )
+            else:
+                tbl_w_meta = plc.io.parquet.read_parquet(
+                    plc.io.SourceInfo(self.paths),
+                    columns=with_columns,
+                    nrows=n_rows,
+                    skip_rows=self.skip_rows,
+                )
+                df = DataFrame.from_table(
+                    tbl_w_meta.tbl,
+                    # TODO: consider nested column names?
+                    tbl_w_meta.column_names(include_children=False),
+                )
+
         elif self.typ == "ndjson":
             json_schema: list[tuple[str, str, list]] = [
                 (name, typ, []) for name, typ in self.schema.items()
@@ -499,13 +539,17 @@ def __init__(self, schema: Schema, key: int, value: IR):
         self.key = key
         self.children = (value,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         try:
             return cache[self.key]
         except KeyError:
             (value,) = self.children
-            return cache.setdefault(self.key, value.evaluate(cache=cache))
+            return cache.setdefault(
+                self.key, value.evaluate(cache=cache, config=config)
+            )
 
 
 class DataFrameScan(IR):
@@ -547,7 +591,9 @@ def get_hashable(self) -> Hashable:
         schema_hash = tuple(self.schema.items())
         return (type(self), schema_hash, id(self.df), self.projection, self.predicate)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         pdf = pl.DataFrame._from_pydf(self.df)
         if self.projection is not None:
@@ -586,10 +632,12 @@ def __init__(
         self.should_broadcast = should_broadcast
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         # Handle any broadcasting
         columns = [e.evaluate(df) for e in self.exprs]
         if self.should_broadcast:
@@ -617,11 +665,11 @@ def __init__(
         self.children = (df,)
 
     def evaluate(
-        self, *, cache: MutableMapping[int, DataFrame]
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
     ) -> DataFrame:  # pragma: no cover; polars doesn't emit this node yet
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         columns = broadcast(*(e.evaluate(df) for e in self.exprs))
         assert all(column.obj.size() == 1 for column in columns)
         return DataFrame(columns)
@@ -700,10 +748,12 @@ def check_agg(agg: expr.Expr) -> int:
         else:
             raise NotImplementedError(f"No handler for {agg=}")
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         keys = broadcast(
             *(k.evaluate(df) for k in self.keys), target_length=df.num_rows
         )
@@ -929,9 +979,11 @@ def _reorder_maps(
             [plc.types.NullOrder.AFTER, plc.types.NullOrder.AFTER],
         ).columns()
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
-        left, right = (c.evaluate(cache=cache) for c in self.children)
+        left, right = (c.evaluate(cache=cache, config=config) for c in self.children)
         how, join_nulls, zlice, suffix, coalesce = self.options
         if how == "cross":
             # Separate implementation, since cross_join returns the
@@ -1038,10 +1090,12 @@ def __init__(
         self.should_broadcast = should_broadcast
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         columns = [c.evaluate(df) for c in self.columns]
         if self.should_broadcast:
             columns = broadcast(*columns, target_length=df.num_rows)
@@ -1095,10 +1149,12 @@ def __init__(
         "any": plc.stream_compaction.DuplicateKeepOption.KEEP_ANY,
     }
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         if self.subset is None:
             indices = list(range(df.num_columns))
             keys_sorted = all(c.is_sorted for c in df.column_map.values())
@@ -1171,10 +1227,12 @@ def __init__(
         self.zlice = zlice
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         sort_keys = broadcast(
             *(k.evaluate(df) for k in self.by), target_length=df.num_rows
         )
@@ -1224,10 +1282,12 @@ def __init__(self, schema: Schema, offset: int, length: int, df: IR):
         self.length = length
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         return df.slice((self.offset, self.length))
 
 
@@ -1244,10 +1304,12 @@ def __init__(self, schema: Schema, mask: expr.NamedExpr, df: IR):
         self.mask = mask
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         (mask,) = broadcast(self.mask.evaluate(df), target_length=df.num_rows)
         return df.filter(mask)
 
@@ -1262,10 +1324,12 @@ def __init__(self, schema: Schema, df: IR):
         self.schema = schema
         self.children = (df,)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
-        df = child.evaluate(cache=cache)
+        df = child.evaluate(cache=cache, config=config)
         # This can reorder things.
         columns = broadcast(
             *(df.column_map[name] for name in self.schema), target_length=df.num_rows
@@ -1333,21 +1397,23 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR):
                 )
             self.options = (tuple(indices), tuple(pivotees), variable_name, value_name)
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (child,) = self.children
         if self.name == "rechunk":
             # No-op in our data model
             # Don't think this appears in a plan tree from python
-            return child.evaluate(cache=cache)  # pragma: no cover
+            return child.evaluate(cache=cache, config=config)  # pragma: no cover
         elif self.name == "rename":
-            df = child.evaluate(cache=cache)
+            df = child.evaluate(cache=cache, config=config)
             # final tag is "swapping" which is useful for the
             # optimiser (it blocks some pushdown operations)
             old, new, _ = self.options
             return df.rename_columns(dict(zip(old, new, strict=True)))
         elif self.name == "explode":
-            df = child.evaluate(cache=cache)
+            df = child.evaluate(cache=cache, config=config)
             ((to_explode,),) = self.options
             index = df.column_names.index(to_explode)
             subset = df.column_names_set - {to_explode}
@@ -1357,7 +1423,7 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         elif self.name == "unpivot":
             indices, pivotees, variable_name, value_name = self.options
             npiv = len(pivotees)
-            df = child.evaluate(cache=cache)
+            df = child.evaluate(cache=cache, config=config)
             index_columns = [
                 Column(col, name=name)
                 for col, name in zip(
@@ -1412,10 +1478,12 @@ def __init__(self, schema: Schema, zlice: tuple[int, int] | None, *children: IR)
         if not all(s.schema == schema for s in self.children[1:]):
             raise NotImplementedError("Schema mismatch")
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         # TODO: only evaluate what we need if we have a slice
-        dfs = [df.evaluate(cache=cache) for df in self.children]
+        dfs = [df.evaluate(cache=cache, config=config) for df in self.children]
         return DataFrame.from_table(
             plc.concatenate.concatenate([df.table for df in dfs]), dfs[0].column_names
         ).slice(self.zlice)
@@ -1459,9 +1527,11 @@ def _extend_with_nulls(table: plc.Table, *, nrows: int) -> plc.Table:
             ]
         )
 
-    def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
+    def evaluate(
+        self, *, cache: MutableMapping[int, DataFrame], config: GPUEngine
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
-        dfs = [df.evaluate(cache=cache) for df in self.children]
+        dfs = [df.evaluate(cache=cache, config=config) for df in self.children]
         max_rows = max(df.num_rows for df in dfs)
         # Horizontal concatenation extends shorter tables with nulls
         dfs = [