[Data] Skip unconditional null strip in find_partition_index

[owenowenisme]

Description

find_partition_index was changed to unconditionally run pd.isna(col_vals) + boolean indexing on every iteration to strip nulls.

This is O(n) with an array allocation on every call, and find_partition_index is called O(blocks × boundaries × columns) times during the sort-shuffle map phase. For SF100 with no nulls (the common case), this added ~10k+ unnecessary allocations per task, causing map tasks to regress from ~3-4 min to ~5-6 min (~50% slower).

Fix: use Arrow's O(1) column.null_count to skip the expensive path when there are no nulls.

Release test result

Regressed runtime: 600+ seconds

Result of case main: {'time': 380.45604542899997, 'object_store_spilled_total_gb': 0.0, 'sf': '100', 'group_by': ['column08', 'column13', 'column14'], 'shuffle_strategy': 'sort_shuffle_pull_based', 'aggregate': True, 'map_groups': False}
--
Finished benchmark, metrics exported to '/tmp/release_test_out.json':
{
"main": {
"time": 424.48728577200006,
"object_store_spilled_total_gb": 0.0,
"sf": "100",
"group_by": [
"column08",
"column13",
"column14"
],
"shuffle_strategy": "sort_shuffle_pull_based",
"aggregate": true,
"map_groups": false
}
}

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 486dfdfbdaca9fd15634dabfb3a7f5c7806ca706. Configure here.}