Implement cudf-polars chunked parquet reading

python/cudf_polars/cudf_polars/callback.py

@@ -18,6 +18,7 @@
 import rmm
 from rmm._cuda import gpu
 
+import cudf_polars.dsl.translate
 from cudf_polars.dsl.translate import translate_ir
 
 if TYPE_CHECKING:
@@ -174,7 +175,9 @@ 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"}):
+    parquet_options = config.config.get("parquet_options", {})
+    cudf_polars.dsl.translate.ir.parquet_options = parquet_options
+    if unsupported := (config.config.keys() - {"raise_on_fail", "parquet_options"}):
         raise ValueError(
             f"Engine configuration contains unsupported settings {unsupported}"
         )

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -35,6 +35,7 @@
 
     from cudf_polars.typing import Schema
 
+parquet_options: dict[str, Any] = {}
 
 __all__ = [
     "IR",
@@ -339,17 +340,46 @@ 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),
-            )
+            if parquet_options.get("chunked", True):
+                reader = plc.io.parquet.ChunkedParquetReader(
+                    plc.io.SourceInfo(self.paths),
+                    columns=with_columns,
+                    num_rows=n_rows,
+                    skip_rows=self.skip_rows,
+                    chunk_read_limit=parquet_options.get("chunk_read_limit", 0),
+                    pass_read_limit=parquet_options.get("pass_read_limit", 1024000000),
+                )
+                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,
+                    num_rows=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()