diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 822a84d38c75..5e5cf8ee12b1 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -606,7 +606,9 @@ def _( rec.state["ir_context"], ch_out, num_producers=num_producers, - estimated_chunk_bytes=executor.target_partition_size, + estimated_chunk_bytes=( + plan.estimated_chunk_bytes or executor.target_partition_size + ), ) ] return nodes, channels diff --git a/python/cudf_polars/cudf_polars/streaming/base.py b/python/cudf_polars/cudf_polars/streaming/base.py index 9415c1d245a1..880393f08619 100644 --- a/python/cudf_polars/cudf_polars/streaming/base.py +++ b/python/cudf_polars/cudf_polars/streaming/base.py @@ -166,12 +166,21 @@ class IOPartitionPlan: - SINGLE_READ: `factor` is the total number of files. """ - __slots__ = ("factor", "flavor") + __slots__ = ("estimated_chunk_bytes", "factor", "flavor") factor: int flavor: IOPartitionFlavor + estimated_chunk_bytes: int | None + """Estimated decompressed bytes per chunk (Scan nodes only).""" - def __init__(self, factor: int, flavor: IOPartitionFlavor) -> None: + def __init__( + self, + factor: int, + flavor: IOPartitionFlavor, + *, + estimated_chunk_bytes: int | None = None, + ) -> None: if flavor == IOPartitionFlavor.SINGLE_FILE and factor != 1: # pragma: no cover raise ValueError(f"Expected factor == 1 for {flavor}, got: {factor}") self.factor = factor self.flavor = flavor + self.estimated_chunk_bytes = estimated_chunk_bytes diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d56c31ce379b..da781fbd8051 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -90,15 +90,19 @@ def scan_partition_plan( if (file_size := sum(column_sizes)) > 0: if file_size > blocksize: # Split large files + factor = math.ceil(file_size / blocksize) return IOPartitionPlan( - math.ceil(file_size / blocksize), + factor, IOPartitionFlavor.SPLIT_FILES, + estimated_chunk_bytes=file_size // factor, ) else: # Fuse small files + factor = max(blocksize // int(file_size), 1) return IOPartitionPlan( - max(blocksize // int(file_size), 1), + factor, IOPartitionFlavor.FUSED_FILES, + estimated_chunk_bytes=file_size * factor, ) # TODO: Use file sizes for csv and json