Use conservative column-size estimates in cudf-polars

python/cudf_polars/cudf_polars/streaming/io.py

@@ -40,7 +40,7 @@
 if TYPE_CHECKING:
     from collections.abc import Hashable, MutableMapping
 
-    from cudf_polars.containers import DataFrame
+    from cudf_polars.containers import DataFrame, DataType
     from cudf_polars.dsl.expr import NamedExpr
     from cudf_polars.dsl.ir import IRExecutionContext
     from cudf_polars.streaming.base import (
@@ -909,6 +909,25 @@ def _sample_rg_sizes(
     return result
 
 
+def _is_fixed_width(dtype: DataType) -> bool:
+    """Return whether dtype is a concrete fixed-width type."""
+    return dtype.id() not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and (
+        plc.traits.is_fixed_width(dtype.plc_type)
+    )
+
+
+def _decoded_size_floor(dtype: DataType, nrows: int) -> int:
+    """Return a conservative decoded-column byte floor for scan planning."""
+    nullmask = (nrows + 7) // 8
+    plc_dtype = dtype.plc_type
+    if dtype.id() == plc.TypeId.STRING:
+        # Decoded strings always have int32 offsets (4 bytes)
+        return (nrows + 1) * 4 + nullmask
+    if _is_fixed_width(dtype):
+        return nrows * plc.types.size_of(plc_dtype) + nullmask
+    return max(1, nrows)
+
+
 class ParquetSourceInfo:
     """Parquet datasource information, fully computed at construction time."""
 
@@ -928,6 +947,7 @@ def from_paths(
         cls,
         paths: tuple[str, ...],
         needed_cols: frozenset[str],
+        schema: tuple[tuple[str, DataType], ...],
         max_footer_samples: int,
         max_row_group_samples: int,
     ) -> ParquetSourceInfo:
@@ -941,37 +961,43 @@ def from_paths(
         if not (file_count and row_count and needed_cols):
             return cls(row_count, {})
 
-        # Floor on size: dictionary encoding can make in-memory size much larger
-        # than what the compressed footer metadata reports.
-        min_floor = max(1, row_count // file_count)
-        suspicious: list[str] = []
+        rows_per_file = max(1, row_count // file_count)
+        schema_map = dict(schema)
+        sample_cols: list[str] = []
 
         for col in needed_cols:
             footer_mean = metadata.mean_size_per_file.get(col)
             if footer_mean is None:
                 continue
-            if footer_mean < min_floor:
-                suspicious.append(col)
+            dtype = schema_map[col]
+            decoded_floor = _decoded_size_floor(dtype, rows_per_file)
+            # This is conservative for all-null columns; footer null counts could
+            # refine the floor later if the extra partitioning becomes costly.
+            if (
+                footer_mean < decoded_floor
+                and max_row_group_samples > 0
+                and not _is_fixed_width(dtype)
+            ):
+                sample_cols.append(col)
             else:
-                per_file_means[col] = footer_mean
+                per_file_means[col] = max(footer_mean, decoded_floor)
 
-        if suspicious and max_row_group_samples > 0:
-            rg_sizes = _sample_rg_sizes(metadata, suspicious, max_row_group_samples)
+        if sample_cols:
+            rg_sizes = _sample_rg_sizes(metadata, sample_cols, max_row_group_samples)
             mean_rg_count = (
                 statistics.mean(metadata.num_row_groups_per_file)
                 if metadata.num_row_groups_per_file
                 else 1
             )
-            for col in suspicious:
+            for col in sample_cols:
                 rg_size = rg_sizes.get(col)
+                decoded_floor = _decoded_size_floor(schema_map[col], rows_per_file)
+                footer_mean = metadata.mean_size_per_file[col]
                 per_file_means[col] = (
-                    max(min_floor, int(rg_size * mean_rg_count))
+                    max(footer_mean, decoded_floor, int(rg_size * mean_rg_count))
                     if rg_size
-                    else min_floor
+                    else max(footer_mean, decoded_floor)
                 )
-        else:
-            for col in suspicious:
-                per_file_means[col] = min_floor
 
         return cls(row_count, per_file_means)
 
@@ -1041,12 +1067,13 @@ def deserialize(cls, data: SerializedDataSourceInfo) -> DataFrameSourceInfo:
 def _build_parquet_source(
     paths: tuple[str, ...],
     needed_cols: frozenset[str],
+    schema: tuple[tuple[str, DataType], ...],
     max_footer_samples: int,
     max_row_group_samples: int,
 ) -> ParquetSourceInfo:
     """Return cached, fully-computed Parquet datasource information."""
     return ParquetSourceInfo.from_paths(
-        paths, needed_cols, max_footer_samples, max_row_group_samples
+        paths, needed_cols, schema, max_footer_samples, max_row_group_samples
     )
 
 
@@ -1055,6 +1082,7 @@ def _build_source_info(
     config_options: ConfigOptions[StreamingExecutor],
     *,
     needed_cols: frozenset[str] | None = None,
+    schema: tuple[tuple[str, DataType], ...] | None = None,
 ) -> DataSourceInfo:
     """Return DataSourceInfo for a Scan or DataFrameScan node."""
     if isinstance(ir, DataFrameScan):
@@ -1063,8 +1091,9 @@ def _build_source_info(
         max_footer = config_options.parquet_options.max_footer_samples
         max_rg = config_options.parquet_options.max_row_group_samples
         needed_cols = frozenset(ir.schema) if needed_cols is None else needed_cols
+        schema = tuple(ir.schema.items()) if schema is None else schema
         paths = tuple(ir.paths)
-        return _build_parquet_source(paths, needed_cols, max_footer, max_rg)
+        return _build_parquet_source(paths, needed_cols, schema, max_footer, max_rg)
     else:  # pragma: no cover
         raise ValueError(f"Unsupported Scan type: {ir.typ}")
 

python/cudf_polars/cudf_polars/streaming/statistics.py

@@ -15,6 +15,7 @@
 
 if TYPE_CHECKING:
     from cudf_polars.dsl.ir import IR
+    from cudf_polars.typing import Schema
     from cudf_polars.utils.config import ConfigOptions, StreamingExecutor
 
 from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars
@@ -41,16 +42,18 @@ def collect_statistics(
     """
     # Group parquet Scan nodes by paths, accumulating the union of needed columns
     # across all Scan nodes that read the same files.
-    parquet_groups: dict[tuple[str, ...], tuple[set[str], list[Scan]]] = {}
+    parquet_groups: dict[tuple[str, ...], tuple[set[str], Schema, list[Scan]]] = {}
     dataframe_scans: list[DataFrameScan] = []
     for node in traversal([root]):
         if isinstance(node, Scan):
             if node.typ == "parquet":
                 paths_key = tuple(node.paths)
                 if paths_key not in parquet_groups:
-                    parquet_groups[paths_key] = (set(), [])
-                parquet_groups[paths_key][0].update(node.schema.keys())
-                parquet_groups[paths_key][1].append(node)
+                    parquet_groups[paths_key] = (set(), {}, [])
+                needed_cols, schema, scan_nodes = parquet_groups[paths_key]
+                needed_cols.update(node.schema.keys())
+                schema.update(node.schema)
+                scan_nodes.append(node)
         elif isinstance(node, DataFrameScan):
             dataframe_scans.append(node)
 
@@ -62,8 +65,9 @@ def collect_statistics(
             scan_nodes[0],
             config_options,
             needed_cols=frozenset(needed_cols),
+            schema=tuple(schema.items()),
         ): scan_nodes
-        for needed_cols, scan_nodes in parquet_groups.values()
+        for needed_cols, schema, scan_nodes in parquet_groups.values()
     }
 
     try:

python/cudf_polars/tests/streaming/test_stats.py

@@ -5,12 +5,15 @@
 
 import json
 import pickle
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar, cast
 
 import pytest
 
 import polars as pl
 
+import pylibcudf as plc
+
+import cudf_polars.streaming.io as streaming_io
 from cudf_polars import Translator
 from cudf_polars.containers import DataType
 from cudf_polars.dsl.ir import Empty, Projection
@@ -121,6 +124,72 @@ def test_base_stats_parquet(
         assert source.column_storage_size("y") is None
 
 
+def test_parquet_source_info_uses_decoded_dtype_floor(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class FakeDataType:
+        def __init__(self, type_id: plc.TypeId) -> None:
+            self.plc_type = plc.DataType(type_id)
+
+        def id(self) -> plc.TypeId:
+            return self.plc_type.id()
+
+    class FakeParquetMetadata:
+        row_count = 2_000
+        mean_size_per_file: ClassVar[dict[str, int]] = {
+            "i64": 1,
+            "dec32": 1,
+            "s": 1,
+            "already_large": 20_000,
+        }
+        num_row_groups_per_file = (1, 1)
+
+        def __init__(self, paths: tuple[str, ...], max_footer_samples: int) -> None:
+            self.paths = paths
+            self.max_footer_samples = max_footer_samples
+
+    sampled_cols: list[str] = []
+
+    def fake_sample_rg_sizes(
+        _metadata: object,
+        target_cols: list[str],
+        _max_row_group_samples: int,
+    ) -> dict[str, int]:
+        sampled_cols.extend(target_cols)
+        return {}
+
+    monkeypatch.setattr(streaming_io, "ParquetMetadata", FakeParquetMetadata)
+    monkeypatch.setattr(streaming_io, "_sample_rg_sizes", fake_sample_rg_sizes)
+
+    source = ParquetSourceInfo.from_paths(
+        ("a.parquet", "b.parquet"),
+        frozenset(
+            {
+                "i64",
+                "dec32",
+                "s",
+                "already_large",
+            }
+        ),
+        (
+            ("i64", DataType(pl.Int64())),
+            ("dec32", cast(DataType, FakeDataType(plc.TypeId.DECIMAL32))),
+            ("s", DataType(pl.String())),
+            ("already_large", DataType(pl.Int64())),
+        ),
+        max_footer_samples=2,
+        max_row_group_samples=1,
+    )
+
+    rows_per_file = 1_000
+    nullmask = 125
+    assert source.column_storage_size("i64") == rows_per_file * 8 + nullmask
+    assert source.column_storage_size("dec32") == rows_per_file * 4 + nullmask
+    assert source.column_storage_size("s") == (rows_per_file + 1) * 4 + nullmask
+    assert source.column_storage_size("already_large") == 20_000
+    assert sampled_cols == ["s"]
+
+
 def test_dataframescan_stats_pickle(
     stats_engine, parquet_stats_executor: concurrent.futures.ThreadPoolExecutor
 ):