snapshot mode in mongodb (#9805)

GitOrigin-RevId: 8f614c5074563e2a59badafd6a6d40df4394ad76

Author: zxqfd555

CHANGELOG.md

@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Added
 - `pw.io.kafka.read` and `pw.io.kafka.write` connectors now support OAUTHBEARER authentication.
+- `pw.io.mongodb.write` connector now supports an `output_table_type` parameter with two modes: `stream_of_changes` (default) and `snapshot`. In `snapshot` mode, the connector maintains the current state of the Pathway table in MongoDB using the `_id` field as the primary key, while `stream_of_changes` preserves the existing behavior by writing all events with `time` and `diff` flags to reflect transactional minibatches and the nature of each change.
+
 
 ## [0.29.0] - 2026-01-22
 

integration_tests/db_connectors/test_mongodb.py

@@ -1,12 +1,54 @@
 import json
+import pathlib
+from typing import Literal
 
-from utils import MONGODB_BASE_NAME, MONGODB_CONNECTION_STRING
+from utils import MONGODB_BASE_NAME, MONGODB_CONNECTION_STRING, MongoDBContext
 
 import pathway as pw
 from pathway.internals.parse_graph import G
 
 
-def test_mongodb(tmp_path, mongodb):
+def write_items_with_connector(
+    *,
+    mongodb: MongoDBContext,
+    test_items: list[dict],
+    input_path: pathlib.Path,
+    schema: type[pw.Schema],
+    output_collection: str,
+    output_table_type: Literal["stream_of_changes", "snapshot"],
+    persistence_config: pw.persistence.Config | None = None,
+) -> list[dict]:
+    G.clear()
+    with open(input_path, "w") as f:
+        for test_item in test_items:
+            f.write(json.dumps(test_item) + "\n")
+    table = pw.io.jsonlines.read(input_path, schema=schema, mode="static")
+    pw.io.mongodb.write(
+        table,
+        connection_string=MONGODB_CONNECTION_STRING,
+        database=MONGODB_BASE_NAME,
+        collection=output_collection,
+        output_table_type=output_table_type,
+    )
+    pw.run(persistence_config=persistence_config)
+
+    result = mongodb.get_collection(output_collection, schema.column_names())
+    result.sort(key=lambda item: (item["name"], item["available"]))
+    return result
+
+
+def check_special_fields(
+    mongodb: MongoDBContext, output_collection: str, *, are_expected: bool
+):
+    full_collection = mongodb.get_full_collection(output_collection)
+    for document in full_collection:
+        time_in_document = "time" in document
+        diff_in_document = "diff" in document
+        assert time_in_document == are_expected, document
+        assert diff_in_document == are_expected, document
+
+
+def test_mongodb_stream_of_changes(tmp_path, mongodb):
     class InputSchema(pw.Schema):
         name: str
         count: int
@@ -15,39 +57,122 @@ class InputSchema(pw.Schema):
 
     input_path = tmp_path / "input.txt"
     output_collection = mongodb.generate_collection_name()
+    test_items = [
+        {"name": "Milk", "count": 500, "price": 1.5, "available": False},
+        {"name": "Water", "count": 600, "price": 0.5, "available": True},
+    ]
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="stream_of_changes",
+    )
+    assert result == test_items
+    check_special_fields(mongodb, output_collection, are_expected=True)
+
+    new_test_items = [{"name": "Milk", "count": 500, "price": 1.5, "available": True}]
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=new_test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="stream_of_changes",
+    )
+    expected_result = [
+        {"name": "Milk", "count": 500, "price": 1.5, "available": False},
+        {"name": "Milk", "count": 500, "price": 1.5, "available": True},
+        {"name": "Water", "count": 600, "price": 0.5, "available": True},
+    ]
+    assert result == expected_result
+    check_special_fields(mongodb, output_collection, are_expected=True)
+
 
-    def run(test_items: list[dict]) -> None:
-        G.clear()
-        with open(input_path, "w") as f:
-            for test_item in test_items:
-                f.write(json.dumps(test_item) + "\n")
-        table = pw.io.jsonlines.read(input_path, schema=InputSchema, mode="static")
-        pw.io.mongodb.write(
-            table,
-            connection_string=MONGODB_CONNECTION_STRING,
-            database=MONGODB_BASE_NAME,
-            collection=output_collection,
-        )
-        pw.run()
+def test_mongodb_snapshot(tmp_path, mongodb):
+    class InputSchema(pw.Schema):
+        name: str = pw.column_definition(primary_key=True)
+        count: int
+        price: float
+        available: bool
+
+    input_path = tmp_path / "input.txt"
+    output_collection = mongodb.generate_collection_name()
 
     test_items = [
         {"name": "Milk", "count": 500, "price": 1.5, "available": False},
         {"name": "Water", "count": 600, "price": 0.5, "available": True},
     ]
-    run(test_items)
-
-    result = mongodb.get_collection(output_collection, InputSchema.column_names())
-    result.sort(key=lambda item: (item["name"], item["available"]))
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="snapshot",
+    )
     assert result == test_items
+    check_special_fields(mongodb, output_collection, are_expected=False)
 
     new_test_items = [{"name": "Milk", "count": 500, "price": 1.5, "available": True}]
-    run(new_test_items)
-
-    result = mongodb.get_collection(output_collection, InputSchema.column_names())
-    result.sort(key=lambda item: (item["name"], item["available"]))
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=new_test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="snapshot",
+    )
     expected_result = [
-        {"name": "Milk", "count": 500, "price": 1.5, "available": False},
         {"name": "Milk", "count": 500, "price": 1.5, "available": True},
         {"name": "Water", "count": 600, "price": 0.5, "available": True},
     ]
     assert result == expected_result
+    check_special_fields(mongodb, output_collection, are_expected=False)
+
+
+def test_mongodb_snapshot_remove(tmp_path, mongodb):
+    class InputSchema(pw.Schema):
+        name: str = pw.column_definition(primary_key=True)
+        count: int
+        price: float
+        available: bool
+
+    input_path = tmp_path / "input.txt"
+    pstorage_path = tmp_path / "PStorage"
+    persistence_config = pw.persistence.Config(
+        backend=pw.persistence.Backend.filesystem(pstorage_path)
+    )
+
+    output_collection = mongodb.generate_collection_name()
+    test_items = [
+        {"name": "Milk", "count": 500, "price": 1.5, "available": False},
+        {"name": "Water", "count": 600, "price": 0.5, "available": True},
+    ]
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="snapshot",
+        persistence_config=persistence_config,
+    )
+    assert result == test_items
+    check_special_fields(mongodb, output_collection, are_expected=False)
+
+    test_items = [
+        {"name": "Water", "count": 600, "price": 0.5, "available": True},
+    ]
+    result = write_items_with_connector(
+        mongodb=mongodb,
+        test_items=test_items,
+        input_path=input_path,
+        schema=InputSchema,
+        output_collection=output_collection,
+        output_table_type="snapshot",
+        persistence_config=persistence_config,
+    )
+    assert result == test_items
+    check_special_fields(mongodb, output_collection, are_expected=False)

integration_tests/db_connectors/utils.py

@@ -256,6 +256,13 @@ def generate_collection_name(self) -> str:
         table_name = f'mongodb_{str(uuid.uuid4()).replace("-", "")}'
         return table_name
 
+    def get_full_collection(
+        self, collection_name: str
+    ) -> list[dict[str, str | int | bool | float]]:
+        db = self.client[MONGODB_BASE_NAME]
+        collection = db[collection_name]
+        return [i for i in collection.find()]  # cast to list
+
     def get_collection(
         self, collection_name: str, field_names: list[str]
     ) -> list[dict[str, str | int | bool | float]]:

python/pathway/io/_utils.py

@@ -28,6 +28,7 @@
 STATIC_MODE_NAME = "static"
 STREAMING_MODE_NAME = "streaming"
 SNAPSHOT_MODE_NAME = "streaming_with_deletions"  # deprecated
+SNAPSHOT_OUTPUT_TABLE_TYPE = "snapshot"
 
 METADATA_COLUMN_NAME = "_metadata"
 MESSAGE_QUEUE_KEY_COLUMN_NAME = "key"

python/pathway/io/deltalake/__init__.py

@@ -25,6 +25,7 @@
 from pathway.internals.table_io import table_from_datasource
 from pathway.internals.trace import trace_user_frame
 from pathway.io._utils import (
+    SNAPSHOT_OUTPUT_TABLE_TYPE,
     _get_unique_name,
     _prepare_s3_connection_settings,
     internal_connector_mode,
@@ -34,7 +35,6 @@
 from pathway.io.s3 import DigitalOceanS3Settings, WasabiS3Settings
 
 _PATHWAY_COLUMN_META_FIELD = "pathway.column.metadata"
-_SNAPSHOT_OUTPUT_TABLE_TYPE = "snapshot"
 _DELTA_LOG_REL_PATH = "_delta_log"
 _LAST_CHECKPOINT_BLOCK_NAME = "_last_checkpoint"
 _CHECKPOINT_EXTENSION = ".checkpoint.parquet"
@@ -618,7 +618,7 @@ def write(
         ),
         min_commit_frequency=min_commit_frequency,
         partition_columns=prepared_partition_columns,
-        snapshot_maintenance_on_output=output_table_type == _SNAPSHOT_OUTPUT_TABLE_TYPE,
+        snapshot_maintenance_on_output=output_table_type == SNAPSHOT_OUTPUT_TABLE_TYPE,
         delta_optimizer_rule=(table_optimizer.engine_rule if table_optimizer else None),
     )
     data_format = api.DataFormat(

python/pathway/io/mongodb/__init__.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
-from typing import Iterable
+from typing import Iterable, Literal
 
 from pathway.internals import api, datasink
 from pathway.internals._io_helpers import _format_output_value_fields
 from pathway.internals.expression import ColumnReference
 from pathway.internals.runtime_type_check import check_arg_types
 from pathway.internals.table import Table
 from pathway.internals.trace import trace_user_frame
+from pathway.io._utils import SNAPSHOT_OUTPUT_TABLE_TYPE
 
 
 @check_arg_types
@@ -20,20 +21,35 @@ def write(
     connection_string: str,
     database: str,
     collection: str,
+    output_table_type: Literal["stream_of_changes", "snapshot"] = "stream_of_changes",
     max_batch_size: int | None = None,
     name: str | None = None,
     sort_by: Iterable[ColumnReference] | None = None,
 ) -> None:
-    """Writes ``table``'s stream of updates to a MongoDB table.
+    """Writes ``table`` to a MongoDB table.
+
+    The output table supports two formats, controlled by the ``output_table_type``
+    parameter.
+
+    The ``stream_of_changes`` format provides a complete history of all modifications
+    applied to the table. Each entry contains the full data row along with two
+    additional fields: ``time`` and ``diff``. The ``time`` field identifies the
+    transactional minibatch in which the change occurred, while ``diff`` describes
+    the nature of the change: ``diff = 1`` indicates that the row was inserted into
+    the Pathway table, and ``diff = -1`` indicates that the row was removed. Row
+    updates are represented as two events within the same transactional minibatch:
+    first the old version of the row with ``diff = -1``, followed by the new version
+    with ``diff = 1``. This format is used by default.
+
+    The ``snapshot`` format maintains the current state of the Pathway table in the
+    output. The table's primary key is stored in the ``_id`` field. When a change
+    occurs, no additional metadata fields are added; instead, the engine locates the
+    corresponding row by ``_id`` and applies the update directly. As a result, the
+    output table always reflects the latest state of the Pathway table.
 
     If the specified database or table doesn't exist, it will be created during the
     first write.
 
-    The entries in the resulting table will have two additional fields: ``time``
-    and ``diff``. In particular, ``time`` is a processing time of a row
-    and ``diff`` shows the nature of the change: ``1`` means a row was added and ``-1``
-    means a row was deleted.
-
     **Note:** Since MongoDB
     `stores DateTime in milliseconds <https://www.mongodb.com/docs/manual/reference/bson-types/#date>`_,
     the `Duration </developers/api-docs/pathway/#pathway.Duration>`_ type is also
@@ -46,6 +62,8 @@ def write(
 for the details.
         database: The name of the database to update.
         collection: The name of the collection to write to.
+        output_table_type: The type of the output table, defining whether a current snapshot
+            or a history of modifications must be maintained.
         max_batch_size: The maximum number of entries to insert in one batch.
         name: A unique name for the connector. If provided, this name will be used in
             logs and monitoring dashboards.
@@ -180,25 +198,68 @@ def write(
     For more advanced setups, such as replica sets, authentication, or custom
     read/write concerns, refer to the official MongoDB documentation on
     `connection strings <https://www.mongodb.com/docs/manual/reference/connection-string/>`_
+
+    Note that if you do not need the full history of modifications, you can use the
+    ``snapshot`` output table type. In this case, the connector configuration would
+    look as follows:
+
+    >>> pw.io.mongodb.write(
+    ...     pet_owners,
+    ...     connection_string="mongodb://127.0.0.1:27017/",
+    ...     database="pathway-test",
+    ...     collection="pet-owners",
+    ...     output_table_type="snapshot",
+    ... )
+
+    The resulting output will be as follows:
+
+    .. code-block:: rst
+
+        [
+            {
+                _id: ObjectId('67180150d94db90697c07853'),
+                age: Long('9'),
+                owner: 'Bob',
+                pet: 'cat',
+            },
+            {
+                _id: ObjectId('67180150d94db90697c07854'),
+                age: Long('8'),
+                owner: 'Alice',
+                pet: 'cat',
+            },
+            {
+                _id: ObjectId('67180150d94db90697c07855'),
+                age: Long('10'),
+                owner: 'Alice',
+                pet: 'dog',
+            }
+        ]
     """
+    is_snapshot_mode = output_table_type == SNAPSHOT_OUTPUT_TABLE_TYPE
     data_storage = api.DataStorage(
         storage_type="mongodb",
         connection_string=connection_string,
         database=database,
         table_name=collection,
         max_batch_size=max_batch_size,
+        snapshot_maintenance_on_output=is_snapshot_mode,
     )
     data_format = api.DataFormat(
         format_type="bson",
         key_field_names=[],
         value_fields=_format_output_value_fields(table),
+        with_special_fields=not is_snapshot_mode,
     )
 
+    datasink_type = (
+        "snapshot" if output_table_type == SNAPSHOT_OUTPUT_TABLE_TYPE else "sink"
+    )
     table.to(
         datasink.GenericDataSink(
             data_storage,
             data_format,
-            datasink_name="mongodb.sink",
+            datasink_name=f"mongodb.{datasink_type}",
             unique_name=name,
             sort_by=sort_by,
         )