wehub-resource-sync
72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
import pandas
|
|
import pytest
|
|
import raydp
|
|
|
|
import ray
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.test_util import _check_usage_record
|
|
|
|
|
|
# RayDP tests require Ray Java. Make sure ray jar is built before running this test.
|
|
@pytest.fixture(scope="function")
|
|
def spark(request):
|
|
ray.init(num_cpus=2, include_dashboard=False)
|
|
spark_session = raydp.init_spark("test", 1, 1, "500M")
|
|
|
|
def stop_all():
|
|
raydp.stop_spark()
|
|
ray.shutdown()
|
|
|
|
request.addfinalizer(stop_all)
|
|
return spark_session
|
|
|
|
|
|
def test_raydp_roundtrip(spark):
|
|
spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
|
|
rows = [(r.one, r.two) for r in spark_df.take(3)]
|
|
ds = ray.data.from_spark(spark_df)
|
|
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
|
assert values == rows
|
|
df = ds.to_spark(spark)
|
|
rows_2 = [(r.one, r.two) for r in df.take(3)]
|
|
assert values == rows_2
|
|
|
|
|
|
def test_raydp_to_spark(spark):
|
|
n = 5
|
|
ds = ray.data.range(n)
|
|
values = [r["id"] for r in ds.take(5)]
|
|
df = ds.to_spark(spark)
|
|
rows = [r.id for r in df.take(5)]
|
|
assert values == rows
|
|
|
|
|
|
def test_from_spark_e2e(spark):
|
|
spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
|
|
|
|
rows = [(r.one, r.two) for r in spark_df.take(3)]
|
|
ds = ray.data.from_spark(spark_df)
|
|
assert len(ds.take_all()) == len(rows)
|
|
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
|
assert values == rows
|
|
|
|
# Check that metadata fetch is included in stats.
|
|
assert "FromArrow" in ds.stats()
|
|
# Underlying implementation uses `FromArrow` operator
|
|
assert ds._logical_plan.dag.name == "FromArrow"
|
|
_check_usage_record(["FromArrow"])
|
|
|
|
|
|
def test_to_pandas(spark):
|
|
df = spark.range(100)
|
|
ds = ray.data.from_spark(df)
|
|
pdf = ds.to_pandas()
|
|
pdf2 = df.toPandas().astype(pdf.dtypes.to_dict())
|
|
pandas.testing.assert_frame_equal(pdf, pdf2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|