chore: import upstream snapshot with attribution

python/ray/data/tests/test_preserve_hash_shuffle_blocks.py

@@ -0,0 +1,73 @@
+import pytest
+
+import ray
+from ray.data.context import DataContext
+from ray.tests.conftest import *  # noqa
+
+
+class TestPreserveHashShuffleBlocks:
+    """Test that hash shuffle repartition preserves block structure."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Setup test context with small target_max_block_size."""
+        ctx = DataContext.get_current()
+        # Very small to force splitting if enabled
+        ctx.target_max_block_size = 1
+        yield
+
+    def test_repartition_preserves_blocks(
+        self, ray_start_regular_shared_2_cpus, shutdown_only
+    ):
+        """Test that repartition with keys preserves block count."""
+        # Create a dataset with multiple blocks
+        ds = ray.data.range(10, override_num_blocks=10)
+
+        # Repartition using hash shuffle with keys
+        result = ds.repartition(2, keys=["id"]).materialize()
+
+        # Should have exactly 2 blocks (one per partition)
+        assert result.num_blocks() == 2
+
+    def test_map_groups_works(self, ray_start_regular_shared_2_cpus, shutdown_only):
+        """Test that map_groups works correctly."""
+        ds = ray.data.from_items(
+            [
+                {"group": 1, "value": 1},
+                {"group": 1, "value": 2},
+                {"group": 2, "value": 3},
+                {"group": 2, "value": 4},
+            ]
+        )
+
+        # map_groups should work correctly
+        result = (
+            ds.groupby("group")
+            .map_groups(lambda g: {"count": [len(g["value"])]})
+            .take_all()
+        )
+
+        assert len(result) == 2
+        counts = sorted([r["count"] for r in result])
+        assert counts == [2, 2]
+
+    def test_join_does_not_preserve_blocks(
+        self, ray_start_regular_shared_2_cpus, shutdown_only
+    ):
+        """Test that join does not preserve block structure (default behavior)."""
+        # Create a dataset with one large block
+        ds = ray.data.range(10, override_num_blocks=2)
+
+        # Join operation uses hash shuffle but doesn't set disallow_block_splitting
+        result = ds.join(
+            ds, on=("id",), join_type="inner", num_partitions=2
+        ).materialize()
+
+        # Should have more than 2 blocks due to block splitting
+        assert result.num_blocks() >= 3
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))