caotranbadat2011 opened a new pull request, #58416:
URL: https://github.com/apache/spark/pull/58416
### What changes were proposed in this pull request?
In `pyspark.pandas`, selecting rows using `.loc[]` with a list-like key
(e.g. `psdf.loc[[key1, key2]]`) delegates to an underlying Spark SQL `.isin()`
filter. Because `.isin()` functions purely as a row filter, keys in the
requested iterable that do not exist in the DataFrame index are silently
discarded, and the matching subset is returned without any warning or error.
This behavior diverges from standard `pandas`, where passing a list of keys
containing any missing key raises a `KeyError` (e.g. `KeyError: '[4] not in
index'`).
An earlier attempt (PR #44236) attempted to resolve this by collecting all
index values via `.tolist()` / `.drop_duplicates().to_pandas()`. However,
collecting the full distributed index onto the driver causes severe network
bottlenecks and leads to Driver Out-Of-Memory (OOM) errors on large datasets.
This PR introduces a distributed set-validation approach with a bounded
memory footprint on the driver:
1. Filters the underlying Spark DataFrame using
`index_col.isin(requested_keys)` and evaluates `.distinct().collect()` solely
on the matched index column.
2. Because only keys present in `requested_keys` can match, the data
transferred to the driver is strictly bounded by $O(K)$, where $K =
|\mathrm{requested\_keys}| \ll N$.
3. Verifies set difference on the driver: `missing_keys = [k for k in
requested_keys if k not in found_keys]`. If `missing_keys` is non-empty, a
`KeyError` is raised, achieving parity with standard pandas behavior.
Fixes [\[SPARK-46306\]](https://issues.apache.org/jira/browse/SPARK-46306).
### Why are the changes needed?
To ensure API consistency and behavioral parity between `pyspark.pandas` and
standard `pandas` when performing label-based indexing with missing keys, while
avoiding driver OOM risks on large distributed datasets.
### Does this PR introduce *any* user-facing change?
Yes.
**Previous behavior:**
```python
psdf = ps.DataFrame({"A": [10, 20, 30]}, index=[1, 2, 3])
psdf.loc[[2, 4]]
# Returns:
# A
# 2 20
# (Key 4 is silently ignored)
```
**New behavior:**
```python
psdf = ps.DataFrame({"A": [10, 20, 30]}, index=[1, 2, 3])
psdf.loc[[2, 4]]
# Raises: KeyError: '[4] not in index'
```
### How was this patch tested?
* Added unit tests in `python/pyspark/pandas/tests/indexes/test_indexing.py`
covering:
* Indexing with fully existing keys.
* Indexing with partially missing keys (asserting `KeyError`).
* Indexing with entirely missing keys (asserting `KeyError`).
* Indexing with a single missing key in a list-like structure `[key]`
(asserting `KeyError`).
* Verified existing indexing test suite passes locally.
### Was this patch authored or co-authored using generative AI tooling?
No.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
