Feature: Rust Based MongoDB Read Connector for Daft

[satya323]

Feature Summary

Add a Rust based daft.read_mongodb() connector that reads MongoDB collections into Daft DataFrames with explicit schema, server-side filter and projection pushdown, and optional distributed range reads. The connector is designed for operational MongoDB workloads: it does not sample or infer schema during planning, and it exposes knobs such as hint, max_time_ms, and user-provided partition_ranges to keep reads bounded and index-friendly. BSON values are converted into the user-provided Daft schema at execution time, with support for nested structs and lists.

Example Usage

import datetime
import daft

schema = {
    "_id": daft.DataType.string(),
    "tenant_id": daft.DataType.string(),
    "created_at": daft.DataType.timestamp("ms"),
    "amount": daft.DataType.float64(),
    "profile": daft.DataType.struct(
        {
            "name": daft.DataType.string(),
            "score": daft.DataType.int64(),
        }
    ),
    "tags": daft.DataType.list(daft.DataType.string()),
}

# Basic read with server-side filter and projection
df = daft.read_mongodb(
    uri="mongodb://localhost:27017/?readPreference=secondaryPreferred",
    database="app",
    collection="orders",
    schema=schema,
    filter={"status": "paid"},
    projection={"_id": 1, "tenant_id": 1, "created_at": 1, "amount": 1, "profile": 1, "tags": 1},
    hint={"status": 1, "created_at": 1},
    max_time_ms=30_000,
    batch_size=1000,
)

# Push additional Daft predicates and projections into MongoDB when supported
result = (
    df.where(daft.col("amount") >= 100)
    .select("tenant_id", "amount", "profile")
)

# Distributed reads across Ray workers using explicit half-open partition ranges
utc = datetime.timezone.utc
df_distributed = daft.read_mongodb(
    uri="mongodb://localhost:27017",
    database="app",
    collection="orders",
    schema=schema,
    filter={"status": "paid"},
    partition_field="created_at",  # may be omitted from output schema
    partition_ranges=[
        (datetime.datetime(2024, 1, 1, tzinfo=utc), datetime.datetime(2024, 2, 1, tzinfo=utc)),
        (datetime.datetime(2024, 2, 1, tzinfo=utc), datetime.datetime(2024, 3, 1, tzinfo=utc)),
        (datetime.datetime(2024, 3, 1, tzinfo=utc), datetime.datetime(2024, 4, 1, tzinfo=utc)),
    ],
    hint={"created_at": 1},
)

# ObjectId and other non-JSON BSON values via MongoDB Extended JSON
df_by_id = daft.read_mongodb(
    uri="mongodb://localhost:27017",
    database="app",
    collection="orders",
    schema=schema,
    filter={"_id": {"$oid": "000000000000000000000000"}},
)

Corner-case Behavior

Schema and BSON conversion

• Explicit schema required: Planning does not inspect the collection. Users must provide the expected Daft schema up front.
• Missing fields / BSON null: Read as Daft nulls.
• Nested documents and arrays: Supported when declared in the schema (e.g. struct, list).
• Type coercion: BSON values are converted to the provided Daft types at execution time; values that cannot be converted should surface as execution errors rather than silently changing semantics.

Filters and predicate pushdown

• User filter + Daft where: Supported MongoDB predicates are merged into a single server-side filter document.
• Supported pushed predicates: Column comparisons (==, !=, <, <=, >, >=), between, is_in with literals, is_null, not_null, and conjunctions of supported predicates.
• Unsupported predicates: Expressions involving casts, functions, arithmetic, not, or unsupported operators are not pushed to MongoDB.
• Strict filter pushdown mode: When daft.context.set_planning_config(enable_strict_filter_pushdown=True) is enabled, unsupported predicates remain above the scan. Without strict mode, planning fails if a pushed filter cannot be fully represented in MongoDB.
• != semantics: Pushed as a null-safe MongoDB filter using $exists, $ne: null, and $ne: <value>.
• Extended JSON in filters: $oid, $numberLong, $numberInt, $numberDouble, $date, and canonical $binary are accepted; malformed or non-canonical extended JSON is rejected at planning time.

Projection pushdown

• Daft select pushdown: Selected columns are pushed into the MongoDB projection when possible.
• Default _id behavior: If _id is not in the Daft schema, MongoDB's default _id field is excluded from the server projection unless the user projection explicitly includes it.
• Projection validation: MongoDB's inclusion/exclusion rules are enforced; mixing inclusion and exclusion fields (except _id) is rejected at planning time.

Distributed reads and partitioning

• partition_ranges omitted: A single scan task is created. No hidden metadata query is issued to infer collection bounds.
• partition_ranges provided: partition_field is required. One MongoDB scan task is created per half-open (lower, upper) range.
• Unbounded ranges: None on the lower or upper bound means unbounded on that side.
• Invalid ranges: Empty or reversed ranges (lower >= upper) are rejected at planning time.
• Partition field outside schema: Allowed when the field is only used to bound MongoDB queries and should not appear in the output DataFrame.
• Range pruning: When filters constrain partition_field with comparable values ($eq, $in, $gt, $gte, $lt, $lte, and supported $and / $or shapes), scan tasks for non-overlapping ranges are skipped.
• Conservative pruning: If a filter contains opaque values or partially unsupported operators on partition_field (e.g. $in with nested documents, mixed $or branches), all partition ranges are retained to avoid dropping data.
• Impossible filters: Contradictory constraints (e.g. {"value": {"$gte": 20, "$lt": 10}} or {"value": {"$in": []}}) result in zero scan tasks.

Limits, batching, and operational safety

• limit pushdown: Applied to MongoDB only when a single scan task remains after partitioning/pruning. For multiple scan tasks, per-task limits are not pushed and Daft applies the limit after the scan.
• max_time_ms: Bounds each server-side find operation.
• hint: Accepts either an index key document or an index name string.
• batch_size: Controls MongoDB cursor batch size and acts as a Daft task batch hint.
• Connection reuse: MongoDB clients are cached per URI within a process.

Out of scope for initial proposal

• Automatic schema inference: Intentionally omitted to avoid expensive collection scans during planning.
• Automatic partition discovery: Users must provide explicit partition_ranges for distributed reads.

[arpatilmh]

This is very useful since nearly 30% of data sources for our data platform are mongodb.

[universalmind303]

@satya323, This would be great as a third party library, but don't think it's right to add this into daft repo.
I suggest implementing this as a custom DataSource: https://docs.daft.ai/en/stable/connectors/custom/#reading-from-a-custom-data-source and publishing it as a third party library if you wish.

[universalmind303]

There's also more information on community extensions here https://docs.daft.ai/en/latest/extensions/community/