Implement cudf-polars chunked parquet reading

python/cudf_polars/cudf_polars/callback.py

@@ -174,7 +174,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}"
         )
@@ -183,7 +183,7 @@ def execute_with_cudf(
             nt.set_udf(
                 partial(
                     _callback,
-                    translate_ir(nt),
+                    translate_ir(nt, config),
                     device=device,
                     memory_resource=memory_resource,
                 )

python/cudf_polars/cudf_polars/dsl/expressions/base.py

@@ -18,6 +18,8 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+    from polars import GPUEngine
+
     from cudf_polars.containers import Column, DataFrame
 
 __all__ = ["Expr", "NamedExpr", "Col", "AggInfo", "ExecutionContext"]
@@ -36,13 +38,16 @@ class ExecutionContext(IntEnum):
 class Expr(Node["Expr"]):
     """An abstract expression object."""
 
-    __slots__ = ("dtype",)
+    __slots__ = ("dtype", "_config")
     dtype: plc.DataType
     """Data type of the expression."""
     # This annotation is needed because of https://github.com/python/mypy/issues/17981
     _non_child: ClassVar[tuple[str, ...]] = ("dtype",)
     """Names of non-child data (not Exprs) for reconstruction."""
 
+    _config: GPUEngine
+    """GPU engine configuration."""
+
     def do_evaluate(
         self,
         df: DataFrame,

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",
@@ -125,12 +129,15 @@ def broadcast(*columns: Column, target_length: int | None = None) -> list[Column
 class IR(Node["IR"]):
     """Abstract plan node, representing an unevaluated dataframe."""
 
-    __slots__ = ("schema",)
+    __slots__ = ("schema", "_config")
     # This annotation is needed because of https://github.com/python/mypy/issues/17981
     _non_child: ClassVar[tuple[str, ...]] = ("schema",)
     schema: Schema
     """Mapping from column names to their data types."""
 
+    _config: GPUEngine
+    """GPU engine configuration."""
+
     def get_hashable(self) -> Hashable:
         """
         Hashable representation of node, treating schema dictionary.
@@ -418,17 +425,51 @@ 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 = self._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
+                    ),
+                )
+                chk = reader.read_chunk()
+                tbl = chk.tbl
+                names = chk.column_names()
+                concatenated_columns = tbl.columns()
+                while reader.has_next():
+                    tbl = reader.read_chunk().tbl
+
+                    for i in range(tbl.num_columns()):
+                        concatenated_columns[i] = plc.concatenate.concatenate(
+                            [concatenated_columns[i], tbl._columns[i]]
+                        )
+                        # Drop residual columns to save memory
+                        tbl._columns[i] = None
+
+                df = DataFrame.from_table(
+                    plc.Table(concatenated_columns),
+                    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()

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -17,16 +17,20 @@
 
 import polars as pl
 import polars.polars as plrs
+from polars import GPUEngine
 from polars.polars import _expr_nodes as pl_expr, _ir_nodes as pl_ir
 
 from cudf_polars.dsl import expr, ir
+from cudf_polars.dsl.expr import Expr
+from cudf_polars.dsl.ir import IR
 from cudf_polars.dsl.traversal import make_recursive, reuse_if_unchanged
 from cudf_polars.typing import NodeTraverser
 from cudf_polars.utils import dtypes, sorting
 
 if TYPE_CHECKING:
     from cudf_polars.typing import ExprTransformer
 
+
 __all__ = ["translate_ir", "translate_named_expr"]
 
 
@@ -361,14 +365,18 @@ def _(
     return ir.HConcat(schema, *(translate_ir(visitor, n=n) for n in node.inputs))
 
 
-def translate_ir(visitor: NodeTraverser, *, n: int | None = None) -> ir.IR:
+def translate_ir(
+    visitor: NodeTraverser, config: GPUEngine = None, *, n: int | None = None
+) -> ir.IR:
     """
     Translate a polars-internal IR node to our representation.
 
     Parameters
     ----------
     visitor
         Polars NodeTraverser object
+    config
+        GPUEngine configuration object
     n
         Optional node to start traversing from, if not provided uses
         current polars-internal node.
@@ -382,6 +390,9 @@ def translate_ir(visitor: NodeTraverser, *, n: int | None = None) -> ir.IR:
     NotImplementedError
         If we can't translate the nodes due to unsupported functionality.
     """
+    if not config:
+        config = GPUEngine()
+
     ctx: AbstractContextManager[None] = (
         set_node(visitor, n) if n is not None else noop_context
     )
@@ -393,6 +404,9 @@ def translate_ir(visitor: NodeTraverser, *, n: int | None = None) -> ir.IR:
             f"No support for polars IR {version=}"
         )  # pragma: no cover; no such version for now.
 
+    IR._config = config
+    Expr._config = config
+
     with ctx:
         polars_schema = visitor.get_schema()
         node = visitor.view_current_node()