chore: import upstream snapshot with attribution

python/ray/data/tests/test_jumbo_arrow_block.py

@@ -0,0 +1,122 @@
+import gc
+import os
+import sys
+from tempfile import TemporaryDirectory
+
+import pyarrow as pa
+import pytest
+from pyarrow import parquet as pq
+
+import ray
+from ray.data._internal.util import GiB, MiB
+from ray.data.context import DataContext
+from ray.tests.conftest import _ray_start
+
+
+@pytest.fixture(scope="module")
+def parquet_dataset_single_column_gt_2gb():
+    chunk_size = 256 * MiB
+    num_chunks = 10
+
+    total_column_size = chunk_size * 10  # ~2.5 GiB
+
+    with TemporaryDirectory() as tmp_dir:
+        dataset_path = f"{tmp_dir}/large_parquet_chunk_{chunk_size}"
+
+        # Create directory
+        os.mkdir(dataset_path)
+
+        for i in range(num_chunks):
+            chunk = b"a" * chunk_size
+
+            d = {"id": [i], "bin": [chunk]}
+            t = pa.Table.from_pydict(d)
+
+            print(f">>> Table schema: {t.schema} (size={sys.getsizeof(t)})")
+
+            filepath = f"{dataset_path}/chunk_{i}.parquet"
+            pq.write_table(t, filepath)
+
+            print(f">>> Created a chunk #{i}")
+
+        print(f">>> Created dataset at {dataset_path}")
+
+        yield dataset_path, num_chunks, total_column_size
+
+        print(f">>> Cleaning up dataset at {dataset_path}")
+
+
+@pytest.fixture(scope="module")
+def ray_cluster_3gb_object_store():
+    original_limit = ray._private.ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT
+
+    ray._private.ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT = 3 * GiB
+
+    with _ray_start(object_store_memory=3 * GiB) as res:
+        yield res
+
+    ray._private.ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT = original_limit
+
+
+@pytest.mark.parametrize(
+    "op",
+    [
+        "map",
+        "map_batches",
+    ],
+)
+@pytest.mark.timeout(300)
+def test_arrow_batch_gt_2gb(
+    ray_cluster_3gb_object_store,
+    parquet_dataset_single_column_gt_2gb,
+    restore_data_context,
+    op,
+):
+    # Disable (automatic) fallback to `ArrowPythonObjectType` extension type
+    DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
+
+    dataset_path, num_rows, total_column_size = parquet_dataset_single_column_gt_2gb
+
+    def _id(x):
+        return x
+
+    ds = ray.data.read_parquet(dataset_path)
+
+    if op == "map":
+        ds = ds.map(_id)
+    elif op == "map_batches":
+        # Combine all rows into a single batch using `map_batches` coercing to
+        # numpy format
+        ds = ds.map_batches(
+            _id,
+            batch_format="pyarrow",
+            batch_size=num_rows,
+            zero_copy_batch=True,
+        )
+
+    batch = ds.take_batch()
+
+    total_binary_column_size = sum([len(b) for b in batch["bin"]])
+
+    print(
+        f">>> Batch:\n"
+        f"------\n"
+        "Column: 'id'\n"
+        f"Values: {batch['id']}\n"
+        f"------\n"
+        "Column: 'bin'\n"
+        f"Total: {total_binary_column_size / GiB} GiB\n"
+        f"Values: {[str(v)[:3] + ' x ' + str(len(v)) for v in batch['bin']]}\n"
+    )
+
+    assert total_binary_column_size == total_column_size
+
+    # Clean up refs
+    del batch
+    del ds
+    # Force GC to free up object store memory
+    gc.collect()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))