⚡️ Speed up function _generate_id_within_group by 138%

[codeflash-ai]

📄 138% (1.38x) speedup for _generate_id_within_group in datacompy/spark/sql.py

⏱️ Runtime : 16.8 seconds → 7.07 seconds (best of 5 runs)

📝 Explanation and details

The optimization achieves a 137% speedup by consolidating multiple expensive Spark operations into a single DataFrame action, eliminating redundant data processing.

Key Optimization: Single-Pass Null/Default Value Detection

The original code performs two separate Spark actions to check for nulls and default values:

null_check = any(list(dataframe.selectExpr(null_cols).first()))      # First Spark action
default_check = any(list(dataframe.selectExpr(default_cols).first()))  # Second Spark action

The optimized version combines both checks into one Spark action:

checks = [(isnull(col(c)), col(c) == lit(default_value)) for c in join_columns]
# Build combined expressions using DataFrame API
first_row = dataframe.select(*check_cols).limit(1).collect()  # Single Spark action

Performance Impact Analysis:

From the line profiler results, the original code spends 95% of total time (26.6s out of 28.2s) on the two separate selectExpr().first() calls. The optimized version reduces this to 86.7% of total time (11.7s out of 13.5s) for the single select().limit(1).collect() operation.

Why This Matters for Hot Path Usage:

Based on the function references, _generate_id_within_group is called during duplicate row processing in _dataframe_merge(), which appears to be part of a core data comparison workflow. When self._any_dupes is True, this function is called twice per comparison (once for each DataFrame), making the optimization particularly valuable for:

• Large datasets with duplicate rows requiring deduplication
• Workflows comparing multiple DataFrames repeatedly
• Data quality pipelines where null handling is frequent

Additional Benefits:

• Reduced DataFrame API overhead: Using native DataFrame functions (isnull, lit) instead of string-based selectExpr reduces parsing overhead
• Better query optimization: Spark's catalyst optimizer can better optimize the combined expression
• Consistent behavior: Single collect operation eliminates potential race conditions between separate checks

The optimization is most effective for datasets with join columns containing nulls or requiring duplicate handling, which are common scenarios in data comparison workflows.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 78 Passed
🌀 Generated Regression Tests	🔘 None Found
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_spark/test_sql_spark.py::test_generate_id_within_group	537ms	175ms	206%✅
test_spark/test_sql_spark.py::test_generate_id_within_group_single_join	1.03s	727ms	41.7%✅

To edit these changes git checkout codeflash/optimize-_generate_id_within_group-mi6f73sp and push.