[damccorm]

Note: this PR was auto-generated by tooling. Before moving it out of draft, the author has done a full review of the PR, specifically evaluating it on the following criteria:

• Is the description accurate? In particular, are the benchmarks correct, meaningful, and reproducible.
• Does the PR avoid doing harm (changes should be non-functional, perf only)?
• (Judgement) Is the code change (and cost of reviewing) worth the performance improvement?

Full details of the methodology of generation and review can be found here - Apache Beam Python SDK Performance PRs Methodology & Branches Impact Ranking

If you have concerns with the PR or with the underlying methodology used to generate this, please comment below - they will be addressed by the author not the tooling :)

Summary of changes:

Optimizes Index.partition_fn (dataframe/partitionings.py) and DataFrameBatchConverter.explode_batch / SeriesBatchConverter (typehints/pandas_type_compatibility.py):

1. Fast-path non-null columns and NamedTuple._make in DataFrameBatchConverter.explode_batch:
   • Uses iter(series) directly when not nulls.any() and constructs rows positionally via map(user_type._make, zip(*iterators)) instead of building and unpacking a {column: value} dict per row.
2. Single-pass df.groupby(pids, sort=False) in Index.partition_fn:
   • Replaces the $O(P \cdot N)$ loop for key in range(num_partitions): yield key, df[hashes % num_partitions == key] (which constructed a boolean mask and scanned the full DataFrame $P$ times) with a single-pass df.groupby(pids, sort=False).

Dataflow benchmarks showing improvement:

• Before the change: 2026-11-12_12_18_07-14933235308757022869
• After the change: 2026-11-12_12_18_08-17650989079603368963
• Summary of perf improvement: In a single-worker (n1-standard-8, --autoscaling_algorithm=NONE --num_workers=1 --max_num_workers=1) Cloud Dataflow pipeline running 30,000,000 schema PerfSchemaRow elements · to_dataframe(pcoll) → filter → to_pcollection(df), built-in transform stage execution time improved from 556,120 ms -> 425,425 ms (1.31x faster, -23.5%) (total single-worker job vCPU time: 1578s -> 1740s).

Additional benchmarks

• Strict Single-PR Isolated Microbenchmark (upstream/master vs upstream/master + this PR alone): bench_50_dataframe_partition_and_explode (c2feb961eef) (on branch perf/py-benchmarks-suite) — 100 iterations of Index.partition_fn(df, 16) + DataFrameBatchConverter.explode_batch(df) on a 5,000-row DataFrame improved from 736.27 ms -> 308.46 ms (2.39x faster, -58.1%, -427.81 ms) in strict isolation.
• Per-Change Attribution (Most to Least Impact, 100 iterations on 5,000-row DataFrame in strict isolation):
  1. Most wall-time impact (83.8% of savings, -322.90 ms): DataFrameBatchConverter.explode_batch using iter(series) on non-null columns and map(user_type._make, ...) — 705.70 ms -> 382.80 ms (1.84x faster, -45.8%).
  2. Second (16.2% of savings, -62.58 ms, scales with partition count $P$): Single-pass df.groupby(pids, sort=False) in Index.partition_fn(df, 16) — 144.90 ms -> 82.32 ms (1.76x faster, -43.2%).
• On-Worker Cloud Dataflow n2-standard-4 Benchmark (p28_ms:28_dataframe_partition_explode_100x5k & suite3_parquet_dataframe): On a Cloud Dataflow n2-standard-4 worker, 100 iterations of partition_fn + explode_batch on 5,000-row DataFrames improved from 1,261 ms -> 529 ms (2.38x faster, -58.0%), and stage_df_partition_wall_ms in suite3_parquet_dataframe improved from 2,678 ms -> 2,210 ms (1.21x faster).

Testing

• Verified with pytest sdks/python/apache_beam/dataframe/partitionings_test.py sdks/python/apache_beam/typehints/pandas_type_compatibility_test.py.

[damccorm]

/gemini review

[damccorm]

/gemini review

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 57.64%. Comparing base (d3078c1) to head (e58ce8d).
⚠️ Report is 111 commits behind head on master.

Files with missing lines	Patch %	Lines
...apache_beam/typehints/pandas_type_compatibility.py	83.33%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #41279      +/-   ##
============================================
+ Coverage     57.48%   57.64%   +0.15%     
  Complexity     2492     2492              
============================================
  Files          1204     1198       -6     
  Lines        171265   171357      +92     
  Branches       2103     2103              
============================================
+ Hits          98456    98772     +316     
+ Misses        70157    69933     -224     
  Partials       2652     2652              

Flag	Coverage Δ
python	80.50% <90.00%> (+0.28%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

📋 Review Summary

This PR cleanly optimizes Index.partition_fn by replacing $O(P \cdot N)$ boolean mask slicing with a single-pass df.groupby(pids, sort=False) and speeds up DataFrameBatchConverter.explode_batch / SeriesBatchConverter by short-circuiting per-element null checks when not nulls.any() and constructing rows positionally via user_type._make / user_type(*values). The changes are well-benchmarked, preserve existing semantics, and have no correctness issues.

🔍 General Feedback

• In Index.partition_fn, using dict(tuple(df.groupby(pids, sort=False))) with an early return for empty frames avoids repeated full-frame scans and works across both pd.DataFrame and pd.Series inputs.
• In DataFrameBatchConverter.explode_batch and SeriesBatchConverter, checking if not nulls.any(): before falling back to element-wise zip(nulls, series) and using map(user_type._make, zip(*iterators)) for NamedTuple row types eliminates per-row dictionary allocation and keyword-argument unpacking overhead.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.