dlt-hub / verified-sources

Contribute to dlt verified sources 🔥
https://dlthub.com/docs/walkthroughs/add-a-verified-source
Apache License 2.0
72 stars 50 forks source link

Add support for pymongoarrow schema and projection parameters in the mongo CollectionLoaders #577

Open marcofraccaro opened 1 month ago

marcofraccaro commented 1 month ago

Source name

mongodb

Describe the data you'd like to see

There are 2 parameters that would be useful to be user-configurable in the mongo CollectionLoaders:

  1. A projection parameter for find_raw_batches/find, which allows to optionally limit which data will be exported from mongo (e.g. to remove at the source columns with sensitive data/reduce data size if not all columns are needed)
  2. A pymongoarrow_schema to be used in PyMongoArrowContext to enforce a schema in process_bson_stream in case it is needed. This means that instead of the current call with the schema set to None as done in context = PyMongoArrowContext.from_schema(None, codec_options=self.collection.codec_options) , one would be able to use a pymongoarrow schema:
            pymongoarrow_schema = pymongoarrow.api.Schema(arrow_schema)
            context = PyMongoArrowContext.from_schema(pymongoarrow_schema, codec_options=self.collection.codec_options)

    Without this schema, in one of our use cases data_item_format = "arrow" fails with the error extraction of resource transaction in generator collection_documents caused an exception: value too large to convert to int32_t. This error is due to the fact that the schema is wrongly inferred to be int32, but setting pyarrow type pa.float64() in the pymongoarrow_schemathings work as expected

Are you a dlt user?

I'm considering using dlt, but this bug is preventing this.

Do you ready to contribute this extension?

Yes, I'm ready.

dlt destination

duckdb/s3

Additional information

No response

esciara commented 4 weeks ago

Careful about using pymongoarrow: we have had a few problems trying to use it, particularly with the translation of ObjectId and arrays of ObjectIds. There is at least the documented problem with nested extension types.

We tried to write pyarrow dataframes using DuckDB's import from Apache Arrow, but it threw an error saying that the type was not supported. We ended up writing pyarrow dataframes straight to parquet files using the pyarrow.parquet.write_table() function, which translates ObjectId to blob (and arrays of blobs respectively), which we then cast using DuckDB's HEX Blob function (which is currently missing from the documentation) to get the id as a string.

marcofraccaro commented 4 weeks ago

@esciara thanks for the heads up! We have indeed noticed similar type-related issues in the past (e.g. https://github.com/mongodb-labs/mongo-arrow/issues/236#issuecomment-2379672931). pymongoarrow is however still very beneficial to us in terms of performances in several use cases.

For ObjectId columns like _id, we are able to use dlt to move data from mongo to duckdb as follows:

  1. We define a pymongoarrow_schema (as explained in the issue description) where ObjectId columns have type pymongoarrow.types.ObjectIdType()
  2. dlt then transforms these columns to string columns with convert_arrow_columns
  3. Duckdb loads these string columns

I have not tried to see what happens with arrays of ObjectIds, but as you noticed this might be tricky.

esciara commented 4 weeks ago

Fab @marcofraccaro. Does it also handle well ObjectId within structures or lists and translates them to string ?

marcofraccaro commented 4 weeks ago

@esciara we have not tried this as it's not needed for our current use case. However based on the limitations of pymongoarrow we both encountered it might not work out of the box.