wehub-resource-sync
214 lines
6.9 KiB
Python
214 lines
6.9 KiB
Python
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__]))
|