chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_csv.py

@@ -0,0 +1,213 @@
+import os
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from packaging.version import Version
+
+import ray
+from ray.data import Schema
+from ray.data._internal.util import rows_same
+from ray.data.block import BlockAccessor
+from ray.data.datasource.path_util import _unwrap_protocol
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+def df_to_csv(dataframe, path, **kwargs):
+    dataframe.to_csv(path, **kwargs)
+
+
+def test_csv_read(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    # Single file.
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.csv")
+    df1.to_csv(path1, index=False)
+    ds = ray.data.read_csv(path1, partitioning=None)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [_unwrap_protocol(path1)]
+    assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
+
+    # Two files, override_num_blocks=2.
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    path2 = os.path.join(tmp_path, "test2.csv")
+    df2.to_csv(path2, index=False)
+    ds = ray.data.read_csv([path1, path2], override_num_blocks=2, partitioning=None)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    df = pd.concat([df1, df2], ignore_index=True)
+    pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    for entry in ds._execute().blocks:
+        assert (
+            # pyrefly: ignore[no-matching-overload]
+            BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
+            == entry.metadata.size_bytes
+        )
+
+    # Three files, override_num_blocks=2.
+    df3 = pd.DataFrame({"one": [7, 8, 9], "two": ["h", "i", "j"]})
+    path3 = os.path.join(tmp_path, "test3.csv")
+    df3.to_csv(path3, index=False)
+    ds = ray.data.read_csv(
+        [path1, path2, path3],
+        override_num_blocks=2,
+        partitioning=None,
+    )
+    df = pd.concat([df1, df2, df3], ignore_index=True)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
+
+
+def test_csv_write(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    input_df = pd.DataFrame({"id": [0]})
+    ds = ray.data.from_blocks([input_df])
+
+    ds.write_csv(tmp_path)
+
+    output_df = pd.concat(
+        [
+            pd.read_csv(os.path.join(tmp_path, filename))
+            for filename in os.listdir(tmp_path)
+        ]
+    )
+    assert rows_same(input_df, output_df)
+
+
+@pytest.mark.parametrize("override_num_blocks", [None, 2])
+def test_csv_roundtrip(
+    ray_start_regular_shared,
+    tmp_path,
+    override_num_blocks,
+    target_max_block_size_infinite_or_default,
+):
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+
+    ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
+    ds.write_csv(tmp_path)
+
+    ds2 = ray.data.read_csv(tmp_path)
+    ds2df = ds2.to_pandas()
+    assert rows_same(ds2df, df)
+    for entry in ds2._execute().blocks:
+        # pyrefly: ignore[no-matching-overload]
+        assert (
+            BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
+            == entry.metadata.size_bytes
+        )
+
+
+def test_csv_read_invalid_format(ray_start_regular_shared, tmp_path):
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+
+    # Setup: CSV and Parquet files in the same directory.
+    csv_path = os.path.join(tmp_path, "test.csv")
+    df.to_csv(csv_path, index=False)
+
+    table = pa.Table.from_pandas(df)
+    parquet_path = os.path.join(tmp_path, "test.parquet")
+    pq.write_table(table, parquet_path)
+
+    # Test 1: CSV parser should fail on Parquet file.
+    error_message = "Failed to read CSV file"
+    with pytest.raises(ValueError, match=error_message):
+        ray.data.read_csv(parquet_path).materialize()
+
+    # Test 2: CSV parser should fail when directory contains non-CSV files.
+    with pytest.raises(ValueError, match=error_message):
+        ray.data.read_csv(tmp_path).materialize()
+
+
+def test_csv_read_no_header(ray_start_regular_shared, tmp_path):
+    from pyarrow import csv
+
+    file_path = os.path.join(tmp_path, "test.csv")
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df.to_csv(file_path, index=False, header=False)
+    ds = ray.data.read_csv(
+        file_path,
+        read_options=csv.ReadOptions(column_names=["one", "two"]),
+    )
+    out_df = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(out_df.dtypes.to_dict()), out_df)
+
+
+def test_csv_read_with_column_type_specified(ray_start_regular_shared, tmp_path):
+    from pyarrow import csv
+
+    file_path = os.path.join(tmp_path, "test.csv")
+    df = pd.DataFrame({"one": [1, 2, 3e1], "two": ["a", "b", "c"]})
+    df.to_csv(file_path, index=False)
+
+    # Incorrect to parse scientific notation in int64 as PyArrow represents
+    # it as double.
+    with pytest.raises(ValueError):
+        ray.data.read_csv(
+            file_path,
+            convert_options=csv.ConvertOptions(
+                column_types={"one": "int64", "two": "string"}
+            ),
+        ).schema()
+
+    # Parsing scientific notation in double should work.
+    ds = ray.data.read_csv(
+        file_path,
+        convert_options=csv.ConvertOptions(
+            column_types={"one": "float64", "two": "string"}
+        ),
+    )
+    expected_df = pd.DataFrame({"one": [1.0, 2.0, 30.0], "two": ["a", "b", "c"]})
+    actual_df = ds.to_pandas()
+    pd.testing.assert_frame_equal(
+        expected_df.astype(actual_df.dtypes.to_dict()), actual_df
+    )
+
+
+@pytest.mark.skipif(
+    Version(pa.__version__) < Version("7.0.0"),
+    reason="invalid_row_handler was added in pyarrow 7.0.0",
+)
+def test_csv_invalid_file_handler(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    from pyarrow import csv
+
+    invalid_txt = "f1,f2\n2,3\nx\n4,5"
+    invalid_file = os.path.join(tmp_path, "invalid.csv")
+    with open(invalid_file, "wt") as f:
+        f.write(invalid_txt)
+
+    ray.data.read_csv(
+        invalid_file,
+        parse_options=csv.ParseOptions(
+            delimiter=",", invalid_row_handler=lambda i: "skip"
+        ),
+    )
+
+
+def test_read_example_data(ray_start_regular_shared, tmp_path):
+    ds = ray.data.read_csv("example://iris.csv")
+    assert ds.count() == 150
+    assert ds.take(1) == [
+        {
+            "sepal.length": 5.1,
+            "sepal.width": 3.5,
+            "petal.length": 1.4,
+            "petal.width": 0.2,
+            "variety": "Setosa",
+        }
+    ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))