diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index e70d28d73863..8892c2d56f30 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 ( @@ -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}") 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..f3fba5cf0bdb 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/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 ):