From f6a65046a9029e32be6b6843e025419223fa46a7 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 23 Jun 2026 11:56:06 +0000 Subject: [PATCH 1/4] Add a I/O partition planning information to benchmark runner --- .../cudf_polars/streaming/benchmarks/utils.py | 25 +++ .../cudf_polars/streaming/explain.py | 208 +++++++++++++++++- .../cudf_polars/cudf_polars/streaming/io.py | 2 +- 3 files changed, 231 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 21b3787ae74c..a2677bd01647 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -263,6 +263,7 @@ class QueryRunResult: plan: SerializablePlan | None iteration_failures: list[tuple[int, int]] validation_failed: bool + partition_plan_rows: list = dataclasses.field(default_factory=list) @dataclasses.dataclass @@ -991,6 +992,16 @@ def run_polars_query( plan = serialize_query(q, engine) + part_plan_rows = [] + if ( + getattr(args, "explain_partition_plan", False) + and engine is not None + and run_config.frontend in _STREAMING_FRONTENDS + ): + from cudf_polars.streaming.explain import collect_partition_plan + + part_plan_rows = collect_partition_plan(q, engine, q_id) + casts = benchmark.EXPECTED_CASTS.get(q_id, []) if numeric_type == "decimal": casts.extend(benchmark.EXPECTED_CASTS_DECIMAL.get(q_id, [])) @@ -1086,6 +1097,7 @@ def run_polars_query( plan=plan, iteration_failures=iteration_failures, validation_failed=validation_failed, + partition_plan_rows=part_plan_rows, ) @@ -1108,6 +1120,7 @@ def _run_query_loop( plans: dict[int, SerializablePlan] = {} validation_failures: list[int] = [] query_failures: list[tuple[int, int]] = [] + all_partition_plan_rows: list = [] for q_id in run_config.queries: try: @@ -1143,6 +1156,12 @@ def _run_query_loop( query_failures.extend(result.iteration_failures) if result.validation_failed: validation_failures.append(q_id) + all_partition_plan_rows.extend(result.partition_plan_rows) + + if all_partition_plan_rows and getattr(args, "explain_partition_plan", False): + from cudf_polars.streaming.explain import format_partition_plan_table + + print(format_partition_plan_table(all_partition_plan_rows), flush=True) return records, plans, validation_failures, query_failures @@ -1942,6 +1961,12 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: help="Print an outline of the logical plan.", default=False, ) + parser.add_argument( + "--explain-partition-plan", + action=argparse.BooleanOptionalAction, + help="Print a combined partition plan summary table across all queries.", + default=False, + ) parser.add_argument( "--print-plans", action=argparse.BooleanOptionalAction, diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 67cdf311ba16..bad29f04b11b 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -9,9 +9,11 @@ import dataclasses import datetime import functools +import os import os.path from collections.abc import Mapping, Sequence from itertools import groupby +from pathlib import Path from typing import TYPE_CHECKING, Any, Self, TypeAlias import cudf_polars.dsl.expressions.binaryop @@ -28,7 +30,7 @@ ) from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal -from cudf_polars.streaming.io import StreamingScan +from cudf_polars.streaming.io import StreamingScan, scan_partition_plan from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.statistics import ( @@ -46,6 +48,20 @@ from cudf_polars.streaming.base import PartitionInfo, StatsCollector +@dataclasses.dataclass +class PartitionPlanRow: + """One row of the partition plan summary table.""" + + query: int + table: str + flavor: str + factor: int + files: int + projected_bytes: int + task_bytes: int + partitions: int + + Serializable: TypeAlias = ( str | int @@ -101,7 +117,7 @@ def explain_query( with cm: stats = collect_statistics(ir, config, executor) lowered_ir, partition_info = lower_ir_graph(ir, config, stats) - return _repr_ir_tree(lowered_ir, partition_info) + return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config) else: if config.executor.name == "streaming": # Include row-count statistics for the logical plan @@ -112,6 +128,171 @@ def explain_query( return _repr_ir_tree(ir) +def collect_partition_plan( + q: pl.LazyFrame, + engine: pl.GPUEngine, + q_id: int, + *, + executor: concurrent.futures.Executor | None = None, +) -> list[PartitionPlanRow]: + """ + Return one PartitionPlanRow per unique StreamingScan in the physical plan. + + Deduplicates scans that appear multiple times due to subquery structure. + """ + cm: contextlib.AbstractContextManager[concurrent.futures.Executor] + if executor is None: + cm = executor = concurrent.futures.ThreadPoolExecutor() + else: + cm = contextlib.nullcontext(executor) + + config = ConfigOptions.from_polars_engine(engine) + ir = Translator(q._ldf.visit(), engine).translate_ir() + + with cm: + stats = collect_statistics(ir, config, executor) + lowered_ir, partition_info = lower_ir_graph(ir, config, stats) + + seen: set[tuple] = set() + rows: list[PartitionPlanRow] = [] + + for node in traversal([lowered_ir]): + if not isinstance(node, StreamingScan): + continue + base_scan = node.base_scan + + dedup_key = (tuple(base_scan.paths), tuple(sorted(base_scan.schema.keys()))) + if dedup_key in seen: + continue + seen.add(dedup_key) + + source = stats.scan_stats.get(base_scan) + if source is None: + continue + + plan = scan_partition_plan(base_scan, stats, config) + projected_bytes = sum( + sz + for col in base_scan.schema + if (sz := source.column_storage_size(col)) is not None + ) + partitions = partition_info[node].count + factor = plan.factor + flavor = plan.flavor.name + + match flavor: + case "SPLIT_FILES": + files = partitions // factor if factor > 0 else partitions + task_bytes = ( + projected_bytes // factor if factor > 0 else projected_bytes + ) + case "FUSED_FILES": + files = partitions * factor + task_bytes = projected_bytes * factor + case _: + files = partitions + task_bytes = projected_bytes + + p = Path(base_scan.paths[0]) + stem = p.stem + parent = p.parent.name + # Prefer the stem unless it looks like a partition filename (purely + # numeric like "1" or prefixed like "part-0"), in which case the + # parent directory holds the table name. + table = parent if (stem.isdigit() or stem.lower().startswith("part")) else stem + + rows.append( + PartitionPlanRow( + query=q_id, + table=table, + flavor=flavor, + factor=factor, + files=files, + projected_bytes=projected_bytes, + task_bytes=task_bytes, + partitions=partitions, + ) + ) + + return rows + + +def _fmt_partition_bytes(b: int) -> str: + if b < 1_000: + return f"{b} B" + elif b < 1_000_000: + return f"{round(b / 1_000, 2):g} KB" + elif b < 1_000_000_000: + return f"{round(b / 1_000_000, 2):g} MB" + else: + return f"{round(b / 1_000_000_000, 2):g} GB" + + +def factor_str(row: PartitionPlanRow) -> str: + """Format the factor field with units appropriate to the scan flavor.""" + match row.flavor: + case "SPLIT_FILES": + return f"{row.factor} tasks/file" + case "FUSED_FILES": + unit = "file" if row.factor == 1 else "files" + return f"{row.factor} {unit}/task" + case _: + return str(row.factor) + + +def format_partition_plan_table(rows: list[PartitionPlanRow]) -> str: + """Format a list of PartitionPlanRows as a fixed-width ASCII table.""" + if not rows: + return "" + + headers = [ + "Q", + "Table", + "Flavor", + "Factor", + "Files", + "Projected (bytes/file)", + "Size/task", + "Partitions", + ] + + flavor_short = {"SPLIT_FILES": "SPLIT", "FUSED_FILES": "FUSED"} + formatted: list[list[str]] = [] + prev_q: int | None = None + for row in rows: + q_str = str(row.query) if row.query != prev_q else "" + prev_q = row.query + formatted.append( + [ + q_str, + row.table, + flavor_short.get(row.flavor, row.flavor), + factor_str(row), + str(row.files), + _fmt_partition_bytes(row.projected_bytes), + _fmt_partition_bytes(row.task_bytes), + str(row.partitions), + ] + ) + + col_widths = [len(h) for h in headers] + for cells in formatted: + for i, cell in enumerate(cells): + col_widths[i] = max(col_widths[i], len(cell)) + + sep = "+-" + "-+-".join("-" * w for w in col_widths) + "-+" + header_row = ( + "| " + " | ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers)) + " |" + ) + lines = ["", "Partition Plan Summary", sep, header_row, sep] + lines.extend( + "| " + " | ".join(c.ljust(col_widths[i]) for i, c in enumerate(cells)) + " |" + for cells in formatted + ) + lines.append(sep) + return "\n".join(lines) + + def serialize_query( q: pl.LazyFrame, engine: pl.GPUEngine, @@ -202,6 +383,7 @@ def _repr_ir_tree( *, offset: str = "", stats: StatsCollector | None = None, + config: ConfigOptions | None = None, ) -> str: header = _repr_ir(ir, offset=offset) count = partition_info[ir].count if partition_info else None @@ -210,11 +392,31 @@ def _repr_ir_tree( row_count_estimate = _fmt_row_count(source.row_count) row_count = f"~{row_count_estimate}" if row_count_estimate else "unknown" header = header.rstrip("\n") + f" {row_count=}\n" + if ( + os.environ.get("CUDF_POLARS__EXPLAIN__PARTITION_PLAN", "0") == "1" + and config is not None + and stats is not None + and isinstance(ir, StreamingScan) + and (source := stats.scan_stats.get(ir.base_scan)) is not None + ): + plan = scan_partition_plan(ir.base_scan, stats, config) + projected_size = sum( + sz + for col in ir.base_scan.schema + if (sz := source.column_storage_size(col)) is not None + ) + plan_info = ( + f"flavor={plan.flavor.name} factor={plan.factor}" + f" projected={_fmt_partition_bytes(projected_size)}" + ) + header = header.rstrip("\n") + f" [{plan_info}]\n" if count is not None: header = header.rstrip("\n") + f" [{count}]\n" children_strs = [ - _repr_ir_tree(child, partition_info, offset=offset + " ", stats=stats) + _repr_ir_tree( + child, partition_info, offset=offset + " ", stats=stats, config=config + ) for child in ir.children ] diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index e70d28d73863..9e135c5e37f6 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -98,7 +98,7 @@ def scan_partition_plan( ) else: # Fuse small files - factor = max(blocksize // int(file_size), 1) + factor = min(max(blocksize // int(file_size), 1), len(ir.paths)) return IOPartitionPlan( factor, IOPartitionFlavor.FUSED_FILES, From 6a8e16e1bd32247df105b36053d62914cfa302fd Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 23 Jun 2026 12:04:58 +0000 Subject: [PATCH 2/4] add tests --- .../tests/streaming/test_explain.py | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 60eb48bb310c..0cc63abae70e 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -15,8 +15,13 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.explain import ( + PartitionPlanRow, + _fmt_partition_bytes, _fmt_row_count, + collect_partition_plan, explain_query, + factor_str, + format_partition_plan_table, serialize_query, ) from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -695,3 +700,167 @@ def test_dynamic_planning_adds_repartition(df, op): assert "REPARTITION" not in plan else: assert "REPARTITION" in plan + + +def test_collect_partition_plan_fused(tmp_path, df): + """Small files with a large target_partition_size → FUSED_FILES.""" + make_partitioned_source(df, tmp_path, fmt="parquet", n_files=4) + q = pl.scan_parquet(tmp_path) + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"target_partition_size": 100_000_000}, + ) + rows = collect_partition_plan(q, engine, q_id=1) + assert len(rows) == 1 + row = rows[0] + assert row.query == 1 + assert row.flavor == "FUSED_FILES" + assert row.factor >= 1 + assert row.files == row.partitions * row.factor + assert row.partitions > 0 + assert row.projected_bytes > 0 + assert row.task_bytes == row.projected_bytes * row.factor + + +def test_collect_partition_plan_split(tmp_path, df): + """Very small target_partition_size → SPLIT_FILES.""" + make_partitioned_source(df, tmp_path, fmt="parquet", n_files=1) + q = pl.scan_parquet(tmp_path) + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"target_partition_size": 1}, + ) + rows = collect_partition_plan(q, engine, q_id=7) + assert len(rows) == 1 + row = rows[0] + assert row.query == 7 + assert row.flavor == "SPLIT_FILES" + assert row.factor > 1 + assert row.files == row.partitions // row.factor + assert row.task_bytes == row.projected_bytes // row.factor + + +def test_collect_partition_plan_table_name_stem(tmp_path, df): + """Single named file: table name is taken from the file stem.""" + single_file = tmp_path / "lineitem.parquet" + df.write_parquet(single_file) + q = pl.scan_parquet(single_file) + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + rows = collect_partition_plan(q, engine, q_id=1) + assert len(rows) == 1 + assert rows[0].table == "lineitem" + + +def test_collect_partition_plan_table_name_parent(tmp_path, df): + """Partitioned directory: table name is taken from the parent directory.""" + (tmp_path / "orders").mkdir() + make_partitioned_source(df, tmp_path / "orders", fmt="parquet", n_files=3) + q = pl.scan_parquet(tmp_path / "orders") + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + rows = collect_partition_plan(q, engine, q_id=1) + assert len(rows) == 1 + assert rows[0].table == "orders" + + +def test_collect_partition_plan_deduplicates(tmp_path, df): + """A scan used twice in a join should produce only one PartitionPlanRow.""" + make_partitioned_source(df, tmp_path, fmt="parquet", n_files=2) + q = pl.scan_parquet(tmp_path) + q = q.join(q, on="x", how="inner") + engine = pl.GPUEngine(executor="streaming", raise_on_fail=True) + rows = collect_partition_plan(q, engine, q_id=1) + assert len(rows) == 1 + + +_SAMPLE_ROWS = [ + PartitionPlanRow( + query=1, + table="lineitem", + flavor="SPLIT_FILES", + factor=2, + files=120, + projected_bytes=2_000_000_000, + task_bytes=1_000_000_000, + partitions=240, + ), + PartitionPlanRow( + query=1, + table="orders", + flavor="FUSED_FILES", + factor=1, + files=60, + projected_bytes=500_000_000, + task_bytes=500_000_000, + partitions=60, + ), + PartitionPlanRow( + query=2, + table="lineitem", + flavor="FUSED_FILES", + factor=1, + files=60, + projected_bytes=1_400_000_000, + task_bytes=1_400_000_000, + partitions=60, + ), +] + + +def test_format_partition_plan_table_empty(): + assert format_partition_plan_table([]) == "" + + +def test_format_partition_plan_table_content(): + table = format_partition_plan_table(_SAMPLE_ROWS) + assert "Partition Plan Summary" in table + assert "lineitem" in table + assert "orders" in table + assert "SPLIT" in table + assert "FUSED" in table + + +def test_format_partition_plan_table_query_shown_once(): + """Each query number should appear on exactly one data row (Q column only).""" + table = format_partition_plan_table(_SAMPLE_ROWS) + # The Q column is the first column; its cell is "| 1 |" for query=1 and "| |" for + # repeated rows. Count occurrences of the literal "| 1 |" at line start. + q1_lines = [line for line in table.splitlines() if line.startswith("| 1 |")] + assert len(q1_lines) == 1 + + +@pytest.mark.parametrize( + "b,expected", + [ + (500, "500 B"), + (1_500, "1.5 KB"), + (2_500_000, "2.5 MB"), + (1_500_000_000, "1.5 GB"), + ], +) +def test_fmt_partition_bytes(b, expected): + assert _fmt_partition_bytes(b) == expected + + +@pytest.mark.parametrize( + "flavor,factor,expected", + [ + ("SPLIT_FILES", 3, "3 tasks/file"), + ("FUSED_FILES", 1, "1 file/task"), + ("FUSED_FILES", 4, "4 files/task"), + ("SINGLE_FILE", 1, "1"), + ], +) +def test_factor_str(flavor, factor, expected): + row = PartitionPlanRow( + query=1, + table="t", + flavor=flavor, + factor=factor, + files=1, + projected_bytes=1, + task_bytes=1, + partitions=1, + ) + assert factor_str(row) == expected From 7129d9cc9d5b7abf9d5cd7eb40e75e8d80e56106 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 23 Jun 2026 13:23:35 +0000 Subject: [PATCH 3/4] copyright --- python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py | 2 +- python/cudf_polars/cudf_polars/streaming/io.py | 2 +- python/cudf_polars/tests/streaming/test_explain.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 2fa92dd2d9dd..42e1722bbc89 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Benchmark utilities for the RapidsMPF SPMD and Ray frontends.""" diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b999196abb16..315e294a21b7 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Multi-partition IO Logic.""" diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5c2b465047b3..79e4db6caf27 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations From 7d20e5d6001aae49161b296c12a9086193d55ece Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 25 Jun 2026 20:47:59 +0000 Subject: [PATCH 4/4] address review --- .../cudf_polars/streaming/explain.py | 26 +++++++------------ .../tests/streaming/test_explain.py | 19 +++++++------- 2 files changed, 19 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 421287a58266..3f6a6dff2869 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -35,6 +35,7 @@ ) from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.base import IOPartitionFlavor from cudf_polars.streaming.io import StreamingScan, scan_partition_plan from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.shuffle import Shuffle @@ -59,7 +60,7 @@ class PartitionPlanRow: query: int table: str - flavor: str + flavor: IOPartitionFlavor factor: int files: int projected_bytes: int @@ -137,24 +138,16 @@ def collect_partition_plan( q: pl.LazyFrame, engine: pl.GPUEngine, q_id: int, - *, - executor: concurrent.futures.Executor | None = None, ) -> list[PartitionPlanRow]: """ Return one PartitionPlanRow per unique StreamingScan in the physical plan. Deduplicates scans that appear multiple times due to subquery structure. """ - cm: contextlib.AbstractContextManager[concurrent.futures.Executor] - if executor is None: - cm = executor = concurrent.futures.ThreadPoolExecutor() - else: - cm = contextlib.nullcontext(executor) - config = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() - with cm: + with concurrent.futures.ThreadPoolExecutor() as executor: stats = collect_statistics(ir, config, executor) lowered_ir, partition_info = lower_ir_graph(ir, config, stats) @@ -183,15 +176,15 @@ def collect_partition_plan( ) partitions = partition_info[node].count factor = plan.factor - flavor = plan.flavor.name + flavor = plan.flavor match flavor: - case "SPLIT_FILES": + case IOPartitionFlavor.SPLIT_FILES: files = partitions // factor if factor > 0 else partitions task_bytes = ( projected_bytes // factor if factor > 0 else projected_bytes ) - case "FUSED_FILES": + case IOPartitionFlavor.FUSED_FILES: files = partitions * factor task_bytes = projected_bytes * factor case _: @@ -236,9 +229,9 @@ def _fmt_partition_bytes(b: int) -> str: def factor_str(row: PartitionPlanRow) -> str: """Format the factor field with units appropriate to the scan flavor.""" match row.flavor: - case "SPLIT_FILES": + case IOPartitionFlavor.SPLIT_FILES: return f"{row.factor} tasks/file" - case "FUSED_FILES": + case IOPartitionFlavor.FUSED_FILES: unit = "file" if row.factor == 1 else "files" return f"{row.factor} {unit}/task" case _: @@ -261,7 +254,6 @@ def format_partition_plan_table(rows: list[PartitionPlanRow]) -> str: "Partitions", ] - flavor_short = {"SPLIT_FILES": "SPLIT", "FUSED_FILES": "FUSED"} formatted: list[list[str]] = [] prev_q: int | None = None for row in rows: @@ -271,7 +263,7 @@ def format_partition_plan_table(rows: list[PartitionPlanRow]) -> str: [ q_str, row.table, - flavor_short.get(row.flavor, row.flavor), + row.flavor.name, factor_str(row), str(row.files), _fmt_partition_bytes(row.projected_bytes), diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 2e68a23b5c9c..5ea6be578ae4 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -19,6 +19,7 @@ from cudf_polars.dsl.expressions.base import Col from cudf_polars.dsl.expressions.binaryop import BinOp from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import IOPartitionFlavor from cudf_polars.streaming.explain import ( PartitionPlanRow, _fmt_partition_bytes, @@ -770,7 +771,7 @@ def test_collect_partition_plan_fused(tmp_path, df): assert len(rows) == 1 row = rows[0] assert row.query == 1 - assert row.flavor == "FUSED_FILES" + assert row.flavor == IOPartitionFlavor.FUSED_FILES assert row.factor >= 1 assert row.files == row.partitions * row.factor assert row.partitions > 0 @@ -791,7 +792,7 @@ def test_collect_partition_plan_split(tmp_path, df): assert len(rows) == 1 row = rows[0] assert row.query == 7 - assert row.flavor == "SPLIT_FILES" + assert row.flavor == IOPartitionFlavor.SPLIT_FILES assert row.factor > 1 assert row.files == row.partitions // row.factor assert row.task_bytes == row.projected_bytes // row.factor @@ -833,7 +834,7 @@ def test_collect_partition_plan_deduplicates(tmp_path, df): PartitionPlanRow( query=1, table="lineitem", - flavor="SPLIT_FILES", + flavor=IOPartitionFlavor.SPLIT_FILES, factor=2, files=120, projected_bytes=2_000_000_000, @@ -843,7 +844,7 @@ def test_collect_partition_plan_deduplicates(tmp_path, df): PartitionPlanRow( query=1, table="orders", - flavor="FUSED_FILES", + flavor=IOPartitionFlavor.FUSED_FILES, factor=1, files=60, projected_bytes=500_000_000, @@ -853,7 +854,7 @@ def test_collect_partition_plan_deduplicates(tmp_path, df): PartitionPlanRow( query=2, table="lineitem", - flavor="FUSED_FILES", + flavor=IOPartitionFlavor.FUSED_FILES, factor=1, files=60, projected_bytes=1_400_000_000, @@ -901,10 +902,10 @@ def test_fmt_partition_bytes(b, expected): @pytest.mark.parametrize( "flavor,factor,expected", [ - ("SPLIT_FILES", 3, "3 tasks/file"), - ("FUSED_FILES", 1, "1 file/task"), - ("FUSED_FILES", 4, "4 files/task"), - ("SINGLE_FILE", 1, "1"), + (IOPartitionFlavor.SPLIT_FILES, 3, "3 tasks/file"), + (IOPartitionFlavor.FUSED_FILES, 1, "1 file/task"), + (IOPartitionFlavor.FUSED_FILES, 4, "4 files/task"), + (IOPartitionFlavor.SINGLE_FILE, 1, "1"), ], ) def test_factor_str(flavor, factor, expected):