chore: import upstream snapshot with attribution

python/ray/data/tests/test_dataset_creation.py

@@ -0,0 +1,68 @@
+import sys
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.execution.interfaces.ref_bundle import (
+    _ref_bundles_iterator_to_block_refs_list,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.util import extract_values
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize(
+    "input_blocks",
+    [
+        [pd.DataFrame({"column": ["spam"]}), pd.DataFrame({"column": ["ham", "eggs"]})],
+        [
+            pa.Table.from_pydict({"column": ["spam"]}),
+            pa.Table.from_pydict({"column": ["ham", "eggs"]}),
+        ],
+    ],
+)
+def test_from_blocks(input_blocks, ray_start_regular_shared):
+    ds = ray.data.from_blocks(input_blocks)
+
+    bundles = ds.iter_internal_ref_bundles()
+    output_blocks = ray.get(_ref_bundles_iterator_to_block_refs_list(bundles))
+    assert len(input_blocks) == len(output_blocks)
+    assert all(
+        input_block.equals(output_block)
+        for input_block, output_block in zip(input_blocks, output_blocks)
+    )
+
+
+def test_from_items(ray_start_regular_shared):
+    ds = ray.data.from_items(["hello", "world"])
+    assert extract_values("item", ds.take()) == ["hello", "world"]
+    assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
+
+
+@pytest.mark.parametrize("parallelism", list(range(1, 21)))
+def test_from_items_parallelism(ray_start_regular_shared, parallelism):
+    # Test that specifying parallelism yields the expected number of blocks.
+    n = 20
+    records = [{"a": i} for i in range(n)]
+    ds = ray.data.from_items(records, override_num_blocks=parallelism)
+    out = ds.take_all()
+    assert out == records
+    assert ds._logical_plan.initial_num_blocks() == parallelism
+
+
+def test_from_items_parallelism_truncated(ray_start_regular_shared):
+    # Test that specifying parallelism greater than the number of items is truncated to
+    # the number of items.
+    n = 10
+    parallelism = 20
+    records = [{"a": i} for i in range(n)]
+    ds = ray.data.from_items(records, override_num_blocks=parallelism)
+    out = ds.take_all()
+    assert out == records
+    assert ds._logical_plan.initial_num_blocks() == n
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))