chore: import upstream snapshot with attribution

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

@@ -0,0 +1,301 @@
+import shutil
+import subprocess
+import tempfile
+import time
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+# To run tests locally, make sure you install mongodb-org and have mongod
+# available on your PATH. Started directly since mongodb-org has no SysV init
+# script. See https://hub.docker.com/_/mongo
+
+
+@pytest.fixture
+def start_mongo():
+    import pymongo
+    import pymongo.errors
+
+    dbpath = tempfile.mkdtemp(prefix="mongod_test_")
+    proc = subprocess.Popen(
+        ["mongod", "--dbpath", dbpath, "--bind_ip", "127.0.0.1"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    # Wait for mongod to accept connections.
+    mongo_url = "mongodb://localhost:27017"
+    for _ in range(30):
+        if proc.poll() is not None:
+            raise RuntimeError(
+                f"mongod exited unexpectedly (returncode={proc.returncode})"
+            )
+        try:
+            client = pymongo.MongoClient(mongo_url, serverSelectionTimeoutMS=1000)
+            client.admin.command("ping")
+            break
+        except pymongo.errors.PyMongoError:
+            time.sleep(0.5)
+    else:
+        proc.kill()
+        raise RuntimeError("mongod failed to start")
+
+    # Make sure a clean slate for each test by dropping
+    # previously created ones (if any).
+    for db in client.list_database_names():
+        # Keep the MongoDB default databases.
+        if db not in ("admin", "local", "config"):
+            client.drop_database(db)
+    yield client, mongo_url
+
+    proc.terminate()
+    proc.wait(timeout=10)
+    shutil.rmtree(dbpath)
+
+
+def test_read_write_mongo(ray_start_regular_shared, start_mongo):
+    from pymongo.errors import ServerSelectionTimeoutError
+    from pymongoarrow.api import Schema
+
+    client, mongo_url = start_mongo
+    foo_db = "foo-db"
+    foo_collection = "foo-collection"
+    foo = client[foo_db][foo_collection]
+    foo.delete_many({})
+
+    # Read nonexistent URI.
+    with pytest.raises(ServerSelectionTimeoutError):
+        ds = ray.data.read_mongo(
+            uri="nonexistent-uri",
+            database=foo_db,
+            collection=foo_collection,
+        )
+    # Read nonexistent database.
+    with pytest.raises(ValueError):
+        ds = ray.data.read_mongo(
+            uri=mongo_url,
+            database="nonexistent-db",
+            collection=foo_collection,
+        )
+    # Read nonexistent collection.
+    with pytest.raises(ValueError):
+        ds = ray.data.read_mongo(
+            uri=mongo_url,
+            database=foo_db,
+            collection="nonexistent-collection",
+        )
+
+    # Inject 5 test docs.
+    docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5)]
+    df = pd.DataFrame(docs).astype({"int_field": "int32"})
+    foo.insert_many(docs)
+
+    # Read a non-empty database, with schema specified.
+    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        schema=schema,
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [3, 2]
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["float_field", "int_field"]
+    assert ds_schema.types == [pa.float64(), pa.int32()]
+    result = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with schema inference, which will read all columns (including the auto
+    # generated internal column "_id").
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read a subset of the collection.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [2, 1]
+    assert ds.count() == 3
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
+
+    # Read with auto-tuned parallelism.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    )
+
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with a parallelism larger than number of rows.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=1000,
+    )
+
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Add a column and then write back to MongoDB.
+    # Inject 2 more test docs.
+    new_docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5, 7)]
+    new_df = pd.DataFrame(new_docs).astype({"int_field": "int32"})
+    ds2 = ray.data.from_pandas(new_df)
+    ds2.write_mongo(uri=mongo_url, database=foo_db, collection=foo_collection)
+
+    # Read again to verify the content.
+    expected_ds = ds.drop_columns(["_id"]).union(ds2)
+    ds3 = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    )
+    ds3.drop_columns(["_id"]).to_pandas().equals(expected_ds.to_pandas())
+
+    # Destination database doesn't exist.
+    with pytest.raises(ValueError):
+        ray.data.range(10).write_mongo(
+            uri=mongo_url, database="nonexistent-db", collection=foo_collection
+        )
+    # Destination collection doesn't exist.
+    with pytest.raises(ValueError):
+        ray.data.range(10).write_mongo(
+            uri=mongo_url, database=foo_db, collection="nonexistent-collection"
+        )
+
+
+def test_mongo_datasource(ray_start_regular_shared, start_mongo):
+    from pymongoarrow.api import Schema
+
+    client, mongo_url = start_mongo
+    foo_db = "foo-db"
+    foo_collection = "foo-collection"
+    foo = client[foo_db][foo_collection]
+    foo.delete_many({})
+
+    # Inject 5 test docs.
+    docs = [{"float_field": 2.0 * key, "int_field": key} for key in range(5)]
+    df = pd.DataFrame(docs).astype({"int_field": "int32"})
+    foo.insert_many(docs)
+
+    # Read non-empty datasource with a specified schema.
+    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        schema=schema,
+        override_num_blocks=2,
+    ).materialize()
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["float_field", "int_field"]
+    assert ds_schema.types == [pa.float64(), pa.int32()]
+    result = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with schema inference, which will read all columns (including the auto
+    # generated internal column "_id").
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=2,
+    ).materialize()
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with auto-tuned parallelism.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    ).materialize()
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with a parallelism larger than number of rows.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=1000,
+    )
+    assert ds.schema(fetch_if_missing=False) is None
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read a subset of the collection.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [2, 1]
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))