From eb4c2d9db814cce21f56a516ded6ea6430153a70 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 5 Jun 2026 07:01:03 -0700 Subject: [PATCH 1/7] add conservative floor logic for column-size estimates --- .../cudf_polars/cudf_polars/streaming/io.py | 75 ++++++++++++++----- .../cudf_polars/streaming/statistics.py | 14 ++-- .../cudf_polars/tests/streaming/test_stats.py | 47 +++++++++++- 3 files changed, 112 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d56c31ce379b..69f0ce0fb779 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/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 ( @@ -813,6 +813,40 @@ def _sample_rg_sizes( return result +def _decoded_size_floor(dtype: DataType, nrows: int) -> int: + """Return a conservative decoded-column byte floor for scan planning.""" + nullmask = (nrows + 7) // 8 + dtype_id = dtype.id() + if dtype_id in (plc.TypeId.INT8, plc.TypeId.UINT8, plc.TypeId.BOOL8): + return nrows + nullmask + if dtype_id in (plc.TypeId.INT16, plc.TypeId.UINT16): + return nrows * 2 + nullmask + if dtype_id in ( + plc.TypeId.INT32, + plc.TypeId.UINT32, + plc.TypeId.FLOAT32, + plc.TypeId.TIMESTAMP_DAYS, + ): + return nrows * 4 + nullmask + if dtype_id in ( + plc.TypeId.INT64, + plc.TypeId.UINT64, + plc.TypeId.FLOAT64, + plc.TypeId.TIMESTAMP_MILLISECONDS, + plc.TypeId.TIMESTAMP_MICROSECONDS, + plc.TypeId.TIMESTAMP_NANOSECONDS, + plc.TypeId.DURATION_MILLISECONDS, + plc.TypeId.DURATION_MICROSECONDS, + plc.TypeId.DURATION_NANOSECONDS, + ): + return nrows * 8 + nullmask + if dtype_id == plc.TypeId.DECIMAL128: + return nrows * 16 + nullmask + if dtype_id == plc.TypeId.STRING: + return (nrows + 1) * 4 + nullmask + return max(1, nrows) + + class ParquetSourceInfo: """Parquet datasource information, fully computed at construction time.""" @@ -832,6 +866,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: @@ -845,37 +880,38 @@ 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) + decoded_floor = _decoded_size_floor(schema_map[col], 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: + 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) @@ -945,12 +981,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 ) @@ -959,6 +996,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): @@ -967,8 +1005,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}") diff --git a/python/cudf_polars/cudf_polars/streaming/statistics.py b/python/cudf_polars/cudf_polars/streaming/statistics.py index ab1e0c1eb51c..30a7fc4eb1d0 100644 --- a/python/cudf_polars/cudf_polars/streaming/statistics.py +++ b/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: diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 36bb7e3fc266..550cd63230a6 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -5,12 +5,13 @@ import json import pickle -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import pytest import polars as pl +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 +122,50 @@ 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 FakeParquetMetadata: + row_count = 2_000 + mean_size_per_file: ClassVar[dict[str, int]] = { + "i64": 1, + "dec": 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 + + def fake_sample_rg_sizes(*args: object) -> dict[str, int]: + 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", "dec", "s", "already_large"}), + ( + ("i64", DataType(pl.Int64())), + ("dec", DataType(pl.Decimal(38, 15))), + ("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("dec") == rows_per_file * 16 + nullmask + assert source.column_storage_size("s") == (rows_per_file + 1) * 4 + nullmask + assert source.column_storage_size("already_large") == 20_000 + + def test_dataframescan_stats_pickle( stats_engine, parquet_stats_executor: concurrent.futures.ThreadPoolExecutor ): From edfbdf86922046ddd187d046721cafcef8989c71 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 5 Jun 2026 08:12:15 -0700 Subject: [PATCH 2/7] address partial code review --- .../cudf_polars/cudf_polars/streaming/io.py | 5 +++ .../cudf_polars/tests/streaming/test_stats.py | 36 +++++++++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 69f0ce0fb779..9766c9c2328d 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -826,18 +826,23 @@ def _decoded_size_floor(dtype: DataType, nrows: int) -> int: plc.TypeId.UINT32, plc.TypeId.FLOAT32, plc.TypeId.TIMESTAMP_DAYS, + plc.TypeId.DURATION_DAYS, + plc.TypeId.DECIMAL32, ): return nrows * 4 + nullmask if dtype_id in ( plc.TypeId.INT64, plc.TypeId.UINT64, plc.TypeId.FLOAT64, + plc.TypeId.TIMESTAMP_SECONDS, plc.TypeId.TIMESTAMP_MILLISECONDS, plc.TypeId.TIMESTAMP_MICROSECONDS, plc.TypeId.TIMESTAMP_NANOSECONDS, + plc.TypeId.DURATION_SECONDS, plc.TypeId.DURATION_MILLISECONDS, plc.TypeId.DURATION_MICROSECONDS, plc.TypeId.DURATION_NANOSECONDS, + plc.TypeId.DECIMAL64, ): return nrows * 8 + nullmask if dtype_id == plc.TypeId.DECIMAL128: diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 550cd63230a6..09ecbef6f1e2 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -5,12 +5,14 @@ import json import pickle -from typing import TYPE_CHECKING, ClassVar +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 @@ -125,10 +127,21 @@ def test_base_stats_parquet( def test_parquet_source_info_uses_decoded_dtype_floor( monkeypatch: pytest.MonkeyPatch, ) -> None: + class FakeDataType: + def __init__(self, type_id: plc.TypeId) -> None: + self.type_id = type_id + + def id(self) -> plc.TypeId: + return self.type_id + class FakeParquetMetadata: row_count = 2_000 mean_size_per_file: ClassVar[dict[str, int]] = { "i64": 1, + "ts": 1, + "duration": 1, + "dec32": 1, + "dec64": 1, "dec": 1, "s": 1, "already_large": 20_000, @@ -147,9 +160,24 @@ def fake_sample_rg_sizes(*args: object) -> dict[str, int]: source = ParquetSourceInfo.from_paths( ("a.parquet", "b.parquet"), - frozenset({"i64", "dec", "s", "already_large"}), + frozenset( + { + "i64", + "ts", + "duration", + "dec32", + "dec64", + "dec", + "s", + "already_large", + } + ), ( ("i64", DataType(pl.Int64())), + ("ts", DataType(pl.Datetime("us"))), + ("duration", DataType(pl.Duration("us"))), + ("dec32", cast(DataType, FakeDataType(plc.TypeId.DECIMAL32))), + ("dec64", cast(DataType, FakeDataType(plc.TypeId.DECIMAL64))), ("dec", DataType(pl.Decimal(38, 15))), ("s", DataType(pl.String())), ("already_large", DataType(pl.Int64())), @@ -161,6 +189,10 @@ def fake_sample_rg_sizes(*args: object) -> dict[str, int]: rows_per_file = 1_000 nullmask = 125 assert source.column_storage_size("i64") == rows_per_file * 8 + nullmask + assert source.column_storage_size("ts") == rows_per_file * 8 + nullmask + assert source.column_storage_size("duration") == rows_per_file * 8 + nullmask + assert source.column_storage_size("dec32") == rows_per_file * 4 + nullmask + assert source.column_storage_size("dec64") == rows_per_file * 8 + nullmask assert source.column_storage_size("dec") == rows_per_file * 16 + nullmask assert source.column_storage_size("s") == (rows_per_file + 1) * 4 + nullmask assert source.column_storage_size("already_large") == 20_000 From c63f5df4b3588790f7f9faf23277689a8a73a9a8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 8 Jun 2026 07:52:44 -0700 Subject: [PATCH 3/7] simplify along tom's suggestion --- .../cudf_polars/cudf_polars/streaming/io.py | 35 +++---------------- .../cudf_polars/tests/streaming/test_stats.py | 4 +-- 2 files changed, 7 insertions(+), 32 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 9766c9c2328d..6a0dd7774f80 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -816,39 +816,14 @@ def _sample_rg_sizes( 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 dtype_id = dtype.id() - if dtype_id in (plc.TypeId.INT8, plc.TypeId.UINT8, plc.TypeId.BOOL8): - return nrows + nullmask - if dtype_id in (plc.TypeId.INT16, plc.TypeId.UINT16): - return nrows * 2 + nullmask - if dtype_id in ( - plc.TypeId.INT32, - plc.TypeId.UINT32, - plc.TypeId.FLOAT32, - plc.TypeId.TIMESTAMP_DAYS, - plc.TypeId.DURATION_DAYS, - plc.TypeId.DECIMAL32, - ): - return nrows * 4 + nullmask - if dtype_id in ( - plc.TypeId.INT64, - plc.TypeId.UINT64, - plc.TypeId.FLOAT64, - plc.TypeId.TIMESTAMP_SECONDS, - plc.TypeId.TIMESTAMP_MILLISECONDS, - plc.TypeId.TIMESTAMP_MICROSECONDS, - plc.TypeId.TIMESTAMP_NANOSECONDS, - plc.TypeId.DURATION_SECONDS, - plc.TypeId.DURATION_MILLISECONDS, - plc.TypeId.DURATION_MICROSECONDS, - plc.TypeId.DURATION_NANOSECONDS, - plc.TypeId.DECIMAL64, - ): - return nrows * 8 + nullmask - if dtype_id == plc.TypeId.DECIMAL128: - return nrows * 16 + nullmask if dtype_id == plc.TypeId.STRING: return (nrows + 1) * 4 + nullmask + if dtype_id not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and ( + plc.traits.is_fixed_width(plc_dtype) + ): + return nrows * plc.types.size_of(plc_dtype) + nullmask return max(1, nrows) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 09ecbef6f1e2..f3d71b8151c2 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -129,10 +129,10 @@ def test_parquet_source_info_uses_decoded_dtype_floor( ) -> None: class FakeDataType: def __init__(self, type_id: plc.TypeId) -> None: - self.type_id = type_id + self.plc_type = plc.DataType(type_id) def id(self) -> plc.TypeId: - return self.type_id + return self.plc_type.id() class FakeParquetMetadata: row_count = 2_000 From d8a800afdb5c486e6e51edd1a947f091dd6155d8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 8 Jun 2026 07:57:19 -0700 Subject: [PATCH 4/7] simplify test --- python/cudf_polars/tests/streaming/test_stats.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index f3d71b8151c2..904063df1c97 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -138,11 +138,7 @@ class FakeParquetMetadata: row_count = 2_000 mean_size_per_file: ClassVar[dict[str, int]] = { "i64": 1, - "ts": 1, - "duration": 1, "dec32": 1, - "dec64": 1, - "dec": 1, "s": 1, "already_large": 20_000, } @@ -163,22 +159,14 @@ def fake_sample_rg_sizes(*args: object) -> dict[str, int]: frozenset( { "i64", - "ts", - "duration", "dec32", - "dec64", - "dec", "s", "already_large", } ), ( ("i64", DataType(pl.Int64())), - ("ts", DataType(pl.Datetime("us"))), - ("duration", DataType(pl.Duration("us"))), ("dec32", cast(DataType, FakeDataType(plc.TypeId.DECIMAL32))), - ("dec64", cast(DataType, FakeDataType(plc.TypeId.DECIMAL64))), - ("dec", DataType(pl.Decimal(38, 15))), ("s", DataType(pl.String())), ("already_large", DataType(pl.Int64())), ), @@ -189,11 +177,7 @@ def fake_sample_rg_sizes(*args: object) -> dict[str, int]: rows_per_file = 1_000 nullmask = 125 assert source.column_storage_size("i64") == rows_per_file * 8 + nullmask - assert source.column_storage_size("ts") == rows_per_file * 8 + nullmask - assert source.column_storage_size("duration") == rows_per_file * 8 + nullmask assert source.column_storage_size("dec32") == rows_per_file * 4 + nullmask - assert source.column_storage_size("dec64") == rows_per_file * 8 + nullmask - assert source.column_storage_size("dec") == rows_per_file * 16 + nullmask assert source.column_storage_size("s") == (rows_per_file + 1) * 4 + nullmask assert source.column_storage_size("already_large") == 20_000 From 0d9c08674f12354097f44b518ca1e1f6e417b98f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 10 Jun 2026 06:03:09 -0700 Subject: [PATCH 5/7] add comment to explain 4 bytes --- python/cudf_polars/cudf_polars/streaming/io.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 6a0dd7774f80..807bb983766a 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -819,6 +819,7 @@ def _decoded_size_floor(dtype: DataType, nrows: int) -> int: plc_dtype = dtype.plc_type dtype_id = dtype.id() if dtype_id == plc.TypeId.STRING: + # Decoded strings always have int32 offsets (4 bytes) return (nrows + 1) * 4 + nullmask if dtype_id not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and ( plc.traits.is_fixed_width(plc_dtype) From 92d0cd8ccfd396d8d6a15166c5e2d810c75b4f84 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 12 Jun 2026 09:12:28 -0700 Subject: [PATCH 6/7] avoid sampling for fixed-width columns --- python/cudf_polars/cudf_polars/streaming/io.py | 17 +++++++++++++++-- .../cudf_polars/tests/streaming/test_stats.py | 10 +++++++++- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 6af7effb19c0..e532baad529e 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -924,6 +924,14 @@ def _decoded_size_floor(dtype: DataType, nrows: int) -> int: return max(1, nrows) +def _has_exact_decoded_size_floor(dtype: DataType) -> bool: + """Return whether the decoded-size floor captures the full data buffer size.""" + dtype_id = dtype.id() + return dtype_id not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and ( + plc.traits.is_fixed_width(dtype.plc_type) + ) + + class ParquetSourceInfo: """Parquet datasource information, fully computed at construction time.""" @@ -965,10 +973,15 @@ def from_paths( footer_mean = metadata.mean_size_per_file.get(col) if footer_mean is None: continue - decoded_floor = _decoded_size_floor(schema_map[col], rows_per_file) + 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: + if ( + footer_mean < decoded_floor + and max_row_group_samples > 0 + and not _has_exact_decoded_size_floor(dtype) + ): sample_cols.append(col) else: per_file_means[col] = max(footer_mean, decoded_floor) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 904063df1c97..f3fba5cf0bdb 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -148,7 +148,14 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int) -> None: self.paths = paths self.max_footer_samples = max_footer_samples - def fake_sample_rg_sizes(*args: object) -> dict[str, int]: + 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) @@ -180,6 +187,7 @@ def fake_sample_rg_sizes(*args: object) -> dict[str, int]: 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( From dd9c2166c35b41a8ea23fee3daba914a557cd638 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 12 Jun 2026 10:23:02 -0700 Subject: [PATCH 7/7] cleanup --- .../cudf_polars/cudf_polars/streaming/io.py | 24 ++++++++----------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index e532baad529e..8892c2d56f30 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -909,29 +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 - dtype_id = dtype.id() - if dtype_id == plc.TypeId.STRING: + if dtype.id() == plc.TypeId.STRING: # Decoded strings always have int32 offsets (4 bytes) return (nrows + 1) * 4 + nullmask - if dtype_id not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and ( - plc.traits.is_fixed_width(plc_dtype) - ): + if _is_fixed_width(dtype): return nrows * plc.types.size_of(plc_dtype) + nullmask return max(1, nrows) -def _has_exact_decoded_size_floor(dtype: DataType) -> bool: - """Return whether the decoded-size floor captures the full data buffer size.""" - dtype_id = dtype.id() - return dtype_id not in (plc.TypeId.EMPTY, plc.TypeId.NUM_TYPE_IDS) and ( - plc.traits.is_fixed_width(dtype.plc_type) - ) - - class ParquetSourceInfo: """Parquet datasource information, fully computed at construction time.""" @@ -980,7 +976,7 @@ def from_paths( if ( footer_mean < decoded_floor and max_row_group_samples > 0 - and not _has_exact_decoded_size_floor(dtype) + and not _is_fixed_width(dtype) ): sample_cols.append(col) else: