feat: Add partition awareness to GroupBy to skip shuffle operations

[Annmool]

Description

This PR adds partition-awareness optimization to Ray Data’s GroupBy operations.

Right now, groupby() always triggers a HashShuffleOperation, even when the dataset is already partitioned by the same column (for example, Hive-style partitioned Parquet data). In those cases, the shuffle is unnecessary and causes avoidable network transfer and memory overhead.

With this change, we first check whether the dataset is already partitioned by the groupby key using block metadata and input file paths. If it is, we skip the shuffle entirely and perform the groupby directly on the existing blocks. If not, the behavior remains unchanged and we fall back to the regular shuffle-based execution.

This optimization is fully backward compatible and does not introduce any API changes.

Changes

1- partition_aware.py (NEW)

• extract_partition_values_from_paths()
  Parses Hive-style file paths and extracts partition column values.

• is_partition_aware_groupby_possible()
  Validates whether the dataset blocks are already partitioned in a way that allows us to safely skip the shuffle phase.

2- grouped_data.py (MODIFIED)

• Added _check_partition_awareness() to the GroupedData class.
• Updated map_groups() to evaluate partition awareness before triggering a shuffle.
• If the dataset is partition-aware, we use it as-is.
• Otherwise, we fall back to the existing shuffle implementation.

3- test_partition_aware_groupby.py (NEW)

• Added unit tests for partition value extraction.
• Added tests for partition-awareness detection logic.
• Covered edge cases such as inconsistent partitions, missing columns, and duplicate partition values.

Impact

• Avoids unnecessary shuffle operations for pre-partitioned datasets.
• Reduces network data movement.
• Improves memory efficiency.
• Provides better performance for Hive-partitioned workloads without changing existing APIs.

[gemini-code-assist]

Code Review

This pull request introduces a valuable optimization for GroupBy operations by adding partition awareness to skip unnecessary shuffles. The implementation, which includes new utility functions for partition detection and modifications to GroupedData, is well-structured. The addition of unit tests is also a positive step.

I have identified a critical issue in how block metadata is collected, which would likely lead to a runtime error. Additionally, I've suggested an improvement to the error handling to enhance observability when the partition awareness check fails. Overall, this is a strong contribution that should significantly improve performance for workloads with pre-partitioned data.

[gemini-code-assist]

The iteration over self._dataset.iter_internal_ref_bundles() appears to be incorrect. This method returns an iterator of RefBundle objects, not tuples. Each RefBundle can contain metadata for multiple blocks.

You should iterate over the RefBundle objects and then access their metadata property to get the list of BlockMetadata for all blocks within that bundle.

            for bundle in self._dataset.iter_internal_ref_bundles():
                blocks_metadata.extend(bundle.metadata)

[gemini-code-assist]

The broad exception handler correctly falls back to the shuffle-based approach, which is a safe default. However, it currently might hide underlying issues in the partition awareness check by only returning the error message as a string.

To improve observability and make debugging easier, I recommend logging the exception as a warning. This will ensure that failures in this optimization path are more visible in the logs.

        except Exception as e:
            # If anything goes wrong, fall back to regular shuffle
            logger = logging.getLogger(__name__)
            logger.warning(
                "Partition awareness check failed, falling back to shuffle: %s",
                e,
                exc_info=True,
            )
            return False, f"Partition awareness check failed: {str(e)}"

[cursor]

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

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[github-actions]

This pull request has been automatically closed because there has been no more activity in the 14 days
since being marked stale.

Please feel free to reopen or open a new pull request if you'd still like this to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for your contribution!