Mongodb source: Arrow loaders broken with pymongoarrow >= 1.6

[mf-unumed]

dlt version

1.17.1

Source name

mongodb

Describe the problem

Summary

CollectionArrowLoader and CollectionArrowLoaderParallel in sources/mongodb/helpers.py rely on pymongoarrow private internals that were removed in pymongoarrow 1.6.0 (released June 2024). Since sources/mongodb/requirements.txt has no upper bound on pymongoarrow, a fresh install today gets 1.14.x and any pipeline using data_item_format="arrow" fails at extraction:

dlt.extract.exceptions.ResourceExtractionError: In processing pipe `<collection>`:
extraction of resource `<collection>` in `generator` `collection_documents` caused an exception:
cannot import name 'process_bson_stream' from 'pymongoarrow.lib'

There are actually two problems; the second one is a silent data-corruption trap for whoever fixes the first.

Problem 1: removed private API

helpers.py#L334-L335 and #L454-L455 do:

from pymongoarrow.context import PyMongoArrowContext
from pymongoarrow.lib import process_bson_stream

In pymongoarrow >= 1.6, process_bson_stream is no longer a module-level function in pymongoarrow.lib (it became a method on PyMongoArrowContext), and PyMongoArrowContext.from_schema(...) was replaced by the plain constructor PyMongoArrowContext(schema, codec_options).

Problem 2: finish() is now cumulative — the loop shape duplicates rows

The current loop (helpers.py#L352-L358) creates one context and calls finish() once per batch:

context = PyMongoArrowContext.from_schema(
    schema=pymongoarrow_schema, codec_options=self.collection.codec_options
)
for batch in cursor:
    process_bson_stream(batch, context)
    table = context.finish()
    yield convert_arrow_columns(table)

On 1.5.2, finish() reset the builders, so each yielded table contained only that batch's rows. Since 1.6, the builders accumulate across finish() calls, so a mechanical port to the new API (just swapping the calls) makes batch N re-yield all previous batches' rows (and in our tests the re-yielded rows come back corrupted as nulls). The fix needs a fresh context per batch, not just new imports.

Reproduction

No MongoDB server needed — BSON batches are fed to pymongoarrow directly, exactly like the raw batches Collection.find_raw_batches() returns, mirroring the load_documents loop:

"""Run with pymongoarrow==1.5.2 and pymongoarrow>=1.6 to compare."""
from importlib.metadata import version

import bson
import pyarrow as pa
from pymongoarrow.api import Schema
from pymongoarrow.context import PyMongoArrowContext

print(f"pymongoarrow {version('pymongoarrow')}, pyarrow {version('pyarrow')}\n")

SCHEMA = Schema({"name": pa.string(), "phone": pa.struct({"e164": pa.string()})})
BATCH_1 = bson.encode({"name": "a", "phone": {"e164": "+1"}})
BATCH_2 = bson.encode({"name": "b", "phone": {"e164": "+2"}})

# 1. The imports used by CollectionArrowLoader (helpers.py L334-335, L454-455)
try:
    from pymongoarrow.lib import process_bson_stream  # noqa: F401
    print("1. `from pymongoarrow.lib import process_bson_stream` -> OK")
except ImportError as exc:
    print(f"1. `from pymongoarrow.lib import process_bson_stream` -> ImportError: {exc}")

has_from_schema = hasattr(PyMongoArrowContext, "from_schema")
print(f"2. PyMongoArrowContext.from_schema exists -> {has_from_schema}")

# 3. The helpers.py loop shape (ONE context, finish() per batch) duplicates rows
#    on >= 1.6, because finish() no longer resets the builders.
if has_from_schema:
    context = PyMongoArrowContext.from_schema(SCHEMA)
else:
    context = PyMongoArrowContext(schema=SCHEMA)

yielded = []
for batch in (BATCH_1, BATCH_2):
    if has_from_schema:
        process_bson_stream(batch, context)
    else:
        context.process_bson_stream(batch)
    table = context.finish()  # helpers.py yields one table per batch here
    yielded.append(table["name"].to_pylist())

print(f"3. per-batch tables: {yielded}")
print("   expected [['a'], ['b']] -- on >=1.6 batch 2 re-yields batch 1's rows\n")

# 4. Bonus (1.6.0 only, fixed in 1.6.4): a batch in which EVERY document is
#    missing a struct-typed schema field crashes finish().
context = (
    PyMongoArrowContext.from_schema(SCHEMA) if has_from_schema else PyMongoArrowContext(schema=SCHEMA)
)
batch_without_phone = bson.encode({"name": "c"})
if has_from_schema:
    process_bson_stream(batch_without_phone, context)
else:
    context.process_bson_stream(batch_without_phone)
try:
    table = context.finish()
    print(f"4. all-missing-struct batch -> OK, phone column: {table['phone'].to_pylist()}")
except ValueError as exc:
    print(f"4. all-missing-struct batch -> ValueError: {exc} (pymongoarrow 1.6.0 only)")

Output on 1.5.2 (current behavior):

1. `from pymongoarrow.lib import process_bson_stream` -> OK
2. PyMongoArrowContext.from_schema exists -> True
3. per-batch tables: [['a'], ['b']]
4. all-missing-struct batch -> OK, phone column: [None]

Output on 1.14.0 (latest):

1. `from pymongoarrow.lib import process_bson_stream` -> ImportError: cannot import name 'process_bson_stream' from 'pymongoarrow.lib'
2. PyMongoArrowContext.from_schema exists -> False
3. per-batch tables: [['a'], [None, 'b']]   <-- batch 2 re-yields batch 1's row, corrupted to null
4. all-missing-struct batch -> OK, phone column: [{'e164': None}]

Suggested fix

We ported our copy of helpers.py like this and it works on 1.6.4+ (both loaders and _run_batch in the parallel loader):

for batch in cursor:
    # finish() no longer resets the builders since pymongoarrow 1.6,
    # so a fresh context is needed per batch
    context = PyMongoArrowContext(
        schema=pymongoarrow_schema, codec_options=self.collection.codec_options
    )
    context.process_bson_stream(batch)
    table = context.finish()
    yield convert_arrow_columns(table)

Another thing worth handling:

• If dropping support for pre-1.6 versions, require pymongoarrow>=1.6.4 rather than >=1.6.0: 1.6.0–1.6.3 have an additional crash — a batch in which every document is missing a struct-typed schema field fails finish() with ValueError: Schema and number of arrays unequal (see repro step 4; fixed upstream by 1.6.4).

Environment

• Python 3.10
• dlt 1.x, pymongo 4.x
• Reproduced with pymongoarrow 1.6.0, 1.6.4, 1.8.0 and 1.14.0 on pyarrow 17/18/19/24

Expected behavior

No response

Steps to reproduce

See description above

How you are using the source?

I run this source in production.

Operating system

Linux

Runtime environment

Local

Python version

3.10

dlt destination

duckdb

Additional information

No response