Turbocharged Data Science: Bridging MongoDB and Python with PyMongoArrow

The world of data is rapidly evolving. The types of data that organizations collect is more voluminous and varied than ever before. At the same time, the tools and techniques for extracting insights from that data are also progressing at breakneck speed.

For data scientists and developers working at this intersection of Big Data and analytics, the ability to efficiently move and shape data between different systems and tools is paramount. Bottlenecks in data transfer can grind pipelines to a halt.

This is where PyMongoArrow comes in. PyMongoArrow is an open source Python library that provides a high-performance bridge between MongoDB, the leading modern general purpose database, and foundational Python data science tools like Pandas, NumPy, and PyArrow.

In this in-depth guide, we‘ll explore what makes PyMongoArrow so powerful and walk through real-world examples of how to use it to accelerate analytics and machine learning workflows. Along the way, we‘ll dive into the technical details and also take a step back to understand where PyMongoArrow fits into the broader data landscape.

Why PyMongoArrow?

Let‘s start with the "why?". After all, there‘s no shortage of methods for moving data between databases and analytics tools. You could use Python libraries like PyMongo to query data from MongoDB and dump it to a Pandas DataFrame or NumPy array directly. Or you could export data from MongoDB to CSV or JSON and then load those into Python. So why introduce yet another tool?

The key lies in performance. Shuffling data between different formats like JSON and CSV is slow and inefficient, especially as data volumes grow. Converting a JSON document to a Pandas DataFrame, for example, requires parsing the nested structure and dynamically allocating memory for the resulting table.

PyMongoArrow solves this by leveraging Apache Arrow. Arrow defines a standardized language-independent memory format for columnar data. PyMongoArrow uses Arrow to efficiently serialize MongoDB‘s BSON document format directly into Arrow, which can then be loaded into Pandas or NumPy with no additional data conversion.

Let‘s look at some benchmarks to underscore the performance difference this makes:

Benchmarks comparing PyMongoArrow to other Python methods for moving data from MongoDB to Pandas. Lower is better. Source: PyMongoArrow documentation

As you can see, PyMongoArrow is nearly 16x faster than using PyMongo to export to Pandas for a 1GB dataset. The speedup is even greater for larger datasets. This is the power of Apache Arrow at work.

But performance is only part of the story. PyMongoArrow also simplifies the code required to move data. The mental overhead of database cursors, schemas, and data serialization is abstracted away. Your Pandas DataFrame or NumPy array is just a method call away from your MongoDB data.

This ease of use makes PyMongoArrow accessible to data scientists who may not be database experts. It also enables developers to build high-performance data pipelines without needing to worry about the intricacies of wrangling data formats.

Connecting to MongoDB

With that background in mind, let‘s dive into some code and see PyMongoArrow in action. We‘ll start by connecting to a MongoDB database.

PyMongoArrow is built on top of PyMongo, the official MongoDB driver for Python. This means we use a standard MongoClient to connect:

from pymongo import MongoClient

client = MongoClient(‘mongodb+srv://user:[email protected]/mydb‘)
db = client.mydb

In this example, we‘re connecting to a MongoDB Atlas cluster. Atlas is MongoDB‘s fully managed database-as-a-service, which makes it easy to spin up a MongoDB database in the cloud with just a few clicks. You can also connect to a self-hosted MongoDB instance running locally or in your datacenter.

With a connection established, let‘s create a MongoDB collection and insert some sample data:

from datetime import datetime

collection = db.mycollection

collection.insert_many([
    {‘_id‘: 1, ‘value‘: 12.34, ‘category‘: ‘A‘, ‘date‘: datetime(2022, 1, 1)},
    {‘_id‘: 2, ‘value‘: 56.78, ‘category‘: ‘B‘, ‘date‘: datetime(2022, 2, 1)},
    {‘_id‘: 3, ‘value‘: 90.12, ‘category‘: ‘C‘, ‘date‘: datetime(2022, 3, 1)},
])

This code creates a new collection called mycollection and inserts three documents into it. Each document has an _id field (which serves as the unique primary key), a value, a category, and a date.

Notice how we‘re using native Python types like float and datetime? This is one of the key advantages of MongoDB‘s document model. Fields can take on any type, and documents in the same collection don‘t have to have the same fields or structure. This flexibility is a major boon for analytics compared to the rigid schemas of traditional SQL databases.

Exporting to Pandas and NumPy

Now that we have some data in MongoDB, let‘s export it to a Pandas DataFrame using PyMongoArrow:

import pandas as pd
from pymongoarrow.monkey import patch_all

patch_all()  # Patch PyMongo to add PyMongoArrow methods
df = collection.find_pandas_all()

print(df)

This code patches the PyMongo Collection object to add PyMongoArrow‘s methods, then calls find_pandas_all() to execute a find query on the collection and load the results into a Pandas DataFrame. The output looks like:

   _id  value category       date
0    1  12.34        A 2022-01-01
1    2  56.78        B 2022-02-01
2    3  90.12        C 2022-03-01

We can just as easily export to a NumPy array using the find_numpy_all() method:

import numpy as np

arr = collection.find_numpy_all()

print(arr)
print(arr.dtype)

Which outputs:

[(1, 12.34, ‘A‘, datetime.datetime(2022, 1, 1, 0, 0))
 (2, 56.78, ‘B‘, datetime.datetime(2022, 2, 1, 0, 0))
 (3, 90.12, ‘C‘, datetime.datetime(2022, 3, 1, 0, 0))]

[(‘_id‘, ‘<i8‘), (‘value‘, ‘<f8‘), (‘category‘, ‘O‘), (‘date‘, ‘O‘)]

The first print statement shows the contents of the NumPy array, and the second shows the dtypes of each field. Notice how the _id field is cast to an int64 while value is a float64.

PyMongoArrow will do its best to infer sensible types, but you can also explicitly specify the schema if you need finer-grained control:

from pymongoarrow.api import Schema

schema = Schema({‘category‘: str, ‘value‘: float})
df = collection.find_pandas_all(schema=schema)

Here we‘ve defined a schema where category is a string and value is a float, and passed it to find_pandas_all(). Any fields not specified in the schema are ignored.

The ability to go from MongoDB to Pandas or NumPy in a single line of code is incredibly powerful. It opens up the full ecosystem of Python data analysis tools, from general purpose libraries like SciPy and scikit-learn to more specialized tools like GeoPandas and TensorFlow.

Two-Way Data Mobility

So far we‘ve focused on exporting data from MongoDB to Python. But what about the reverse? PyMongoArrow also supports writing Pandas DataFrames, NumPy arrays, and Arrow tables back to MongoDB:

from pymongoarrow.api import write

df = pd.DataFrame({‘category‘: [‘A‘, ‘B‘], ‘value‘: [12.34, 56.78]})
write(db.mynewcollection, df)

This creates a new collection called mynewcollection and writes the contents of the DataFrame to it. The collection is created if it doesn‘t already exist.

Being able to round-trip data between systems like this is essential for machine learning workflows. You can pull data into Python for model training and evaluation and then write results and predictions back to MongoDB for downstream applications to consume. PyMongoArrow makes this trivial.

Faster ETL Pipelines

Another major use case for PyMongoArrow is building high-performance ETL (extract, transform, load) pipelines. The Arrow memory format is supported by a wide range of data systems beyond just Python – including Spark, Dask, Kafka, and more. This allows you to build end-to-end data pipelines that avoid costly serialization overhead.

For instance, you could use PyMongoArrow to extract data from MongoDB and load it into Spark for distributed processing using PySpark. From there you could write the transformed data back to MongoDB or send it to another system that supports Arrow, like Amazon S3.

Alternatively, you could build a streaming pipeline where MongoDB Change Streams emit data changes as Arrow records that get picked up by a Kafka Connect sink and published to Kafka topics. A Spark Structured Streaming job could then subscribe to those topics, process the data, and write it to MongoDB or another database.

By keeping data in the Arrow format throughout, you eliminate serialization bottlenecks and enable seamless interoperability between systems.

An Extensible Ecosystem

MongoDB has a rich ecosystem of tools and integrations beyond those we‘ve covered in this guide. The MongoDB Connector for Apache Spark allows you to use MongoDB as a source and sink for Spark jobs. MongoDB Charts makes it easy to visually explore and dashboard your MongoDB data. MongoDB Atlas Data Lake enables S3 querying and analytics on your data while it‘s still in MongoDB. And the list goes on.

As more tools add Arrow support, PyMongoArrow will serve as the glue that connects them to MongoDB. The recently released Siuba (https://github.com/machow/siuba) library, for instance, provides an API for querying Arrow tables with a SQL-like syntax. PyMongoArrow allows you to get a Siuba table representation of your MongoDB collection in one line:

import siuba

tbl = collection.find_arrow_all().to_siuba()
tbl.filter(tbl.value > 50).collect()

This is just a small taste of what‘s possible. As the Arrow ecosystem grows, so too will the possibilities for integrating MongoDB with other systems using PyMongoArrow.

Looking Forward

As data volumes continue to grow and machine learning becomes more widely adopted, the ability to efficiently move and process data between different systems will only become more critical. We believe PyMongoArrow is well-positioned to serve this need and bridge the gap between operational databases like MongoDB and the tools data scientists and engineers use every day.

We‘re excited to see what the community builds with PyMongoArrow. In the short time since its release, it has already been incorporated into downstream libraries like the MongoML machine learning tool (https://github.com/rlondner/mongoml).

If you‘re interested in learning more, check out the PyMongoArrow documentation (https://mongo-arrow.readthedocs.io/en/latest/) and the MongoDB Developer Hub (https://developer.mongodb.com) for tutorials and guides. Feel free to join the discussion on the MongoDB Community Forums (https://developer.mongodb.com/community/forums/) and let us know how you‘re using PyMongoArrow.

Happy coding, and may your analytics be speedy!