chore: import upstream snapshot with attribution

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

@@ -0,0 +1,71 @@
+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__]))