code: 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,
          parallelism=2,
      )
      assert ds._block_num_rows() == [3, 2]
      assert str(ds) == (
          "Dataset(\n"
          "   num_blocks=2,\n"
          "   num_rows=5,\n"
          "   schema={float_field: double, int_field: int32}\n"
          ")"
      )
      assert df.equals(ds.to_pandas())
  
      # 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,
          parallelism=2,
      )
      assert ds._block_num_rows() == [3, 2]
      assert str(ds) == (
          "Dataset(\n"
          "   num_blocks=2,\n"
          "   num_rows=5,\n"
          "   schema={_id: fixed_size_binary[12], float_field: double, "
          "int_field: int32}\n"
          ")"
      )
      assert df.equals(ds.drop_columns(["_id"]).to_pandas())
  
      # 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}}}],
          parallelism=2,
      )
      assert ds._block_num_rows() == [2, 1]
      assert str(ds) == (
          "Dataset(\n"
          "   num_blocks=2,\n"
          "   num_rows=3,\n"
          "   schema={_id: fixed_size_binary[12], float_field: double, "
          "int_field: int32}\n"
          ")"
      )
      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 str(ds) == (
          "Dataset(\n"
          "   num_blocks=200,\n"
          "   num_rows=5,\n"
          "   schema={_id: fixed_size_binary[12], float_field: double, "
          "int_field: int32}\n"
          ")"
      )
      assert df.equals(ds.drop_columns(["_id"]).to_pandas())
  
      # Read with a parallelism larger than number of rows.
      ds = ray.data.read_mongo(
          uri=mongo_url,
          database=foo_db,
          collection=foo_collection,
          parallelism=1000,
      )
      assert str(ds) == (
          "Dataset(\n"
          "   num_blocks=1000,\n"
          "   num_rows=5,\n"
          "   schema={_id: fixed_size_binary[12], float_field: double, "
          "int_field: int32}\n"
          ")"
      )
      assert df.equals(ds.drop_columns(["_id"]).to_pandas())
  
      # 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",
          )
  
file_path: ray/python/ray/data/tests/test_mongo.py
imports: import subprocess
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data.tests.conftest import *
  from ray.data.tests.mock_http_server import *
  from ray.tests.conftest import *
label_desc: Use image node_id 2 for calling a global function with example usage: test_read_write_mongo(ray_start_regular_shared, start_mongo) without return types