chore: import upstream snapshot with attribution

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

@@ -0,0 +1,161 @@
+import os
+import zipfile
+
+import pytest
+from packaging.version import parse as parse_version
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.datasource.path_util import (
+    _resolve_paths_and_filesystem,
+    _unwrap_protocol,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+MIN_PYARROW_VERSION_FOR_HUDI = parse_version("11.0.0")
+PYARROW_VERSION = get_pyarrow_version()
+PYARROW_VERSION_MEETS_REQUIREMENT = (
+    PYARROW_VERSION and PYARROW_VERSION >= MIN_PYARROW_VERSION_FOR_HUDI
+)
+PYARROW_HUDI_TEST_SKIP_REASON = (
+    f"Hudi only supported if pyarrow >= {MIN_PYARROW_VERSION_FOR_HUDI}"
+)
+
+
+def _extract_testing_table(fixture_path: str, table_dir: str, target_dir: str) -> str:
+    with zipfile.ZipFile(fixture_path, "r") as zip_ref:
+        zip_ref.extractall(target_dir)
+    return os.path.join(target_dir, table_dir)
+
+
+def _get_hudi_table_path(fs, data_path, table_name, testing_dir="test_hudi") -> str:
+    setup_data_path = _unwrap_protocol(data_path)
+    target_testing_dir = os.path.join(setup_data_path, testing_dir)
+    fixture_path, _ = _resolve_paths_and_filesystem(
+        f"example://hudi-tables/{table_name}.zip", fs
+    )
+    return _extract_testing_table(fixture_path[0], table_name, target_testing_dir)
+
+
+@pytest.mark.skipif(
+    not PYARROW_VERSION_MEETS_REQUIREMENT,
+    reason=PYARROW_HUDI_TEST_SKIP_REASON,
+)
+@pytest.mark.parametrize(
+    "fs,data_path",
+    [
+        (None, lazy_fixture("local_path")),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
+    ],
+)
+def test_hudi_snapshot_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
+    table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
+
+    ds = ray.data.read_hudi(table_path, filters=[("city", "=", "san_francisco")])
+
+    assert ds.schema().names == [
+        "_hoodie_commit_time",
+        "_hoodie_commit_seqno",
+        "_hoodie_record_key",
+        "_hoodie_partition_path",
+        "_hoodie_file_name",
+        "ts",
+        "uuid",
+        "rider",
+        "driver",
+        "fare",
+        "city",
+    ]
+    assert ds.count() == 4
+    rows = (
+        ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"])
+        .sort("fare")
+        .take_all()
+    )
+    first_commit = "20250715043008154"
+    second_commit = "20250715043011090"
+    assert rows == [
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695159649087,
+            "rider": "rider-A",
+            "fare": 19.10,
+        },
+        {
+            "_hoodie_commit_time": second_commit,
+            "ts": 1695046462179,
+            "rider": "rider-D",
+            "fare": 25.0,
+        },
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695091554788,
+            "rider": "rider-C",
+            "fare": 27.70,
+        },
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695332066204,
+            "rider": "rider-E",
+            "fare": 93.50,
+        },
+    ]
+
+
+@pytest.mark.skipif(
+    not PYARROW_VERSION_MEETS_REQUIREMENT,
+    reason=PYARROW_HUDI_TEST_SKIP_REASON,
+)
+@pytest.mark.parametrize(
+    "fs,data_path",
+    [
+        (None, lazy_fixture("local_path")),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
+    ],
+)
+def test_hudi_incremental_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
+    table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
+
+    first_commit = "20250715043008154"
+    second_commit = "20250715043011090"
+    ds = ray.data.read_hudi(
+        table_path,
+        query_type="incremental",
+        hudi_options={
+            "hoodie.read.file_group.start_timestamp": first_commit,
+            "hoodie.read.file_group.end_timestamp": second_commit,
+        },
+    )
+
+    assert ds.schema().names == [
+        "_hoodie_commit_time",
+        "_hoodie_commit_seqno",
+        "_hoodie_record_key",
+        "_hoodie_partition_path",
+        "_hoodie_file_name",
+        "ts",
+        "uuid",
+        "rider",
+        "driver",
+        "fare",
+        "city",
+    ]
+    assert ds.count() == 1
+    rows = ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"]).take_all()
+    assert rows == [
+        {
+            "_hoodie_commit_time": second_commit,
+            "ts": 1695046462179,
+            "rider": "rider-D",
+            "fare": 25.0,
+        },
+    ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))