Inicio rápido

Requisitos previos

Asegúrese de tener la distribución PyMongoArrow Instalado. En el shell de Python, lo siguiente debería ejecutarse sin generar una excepción:

>>> import pymongoarrow as pma

Este tutorial también asume que una instancia de MongoDB se ejecuta en el host y puerto predeterminados. Después de descargar e instalar MongoDB, puede iniciarlo como se muestra en el siguiente ejemplo de código:

$ mongod

from pymongoarrow.monkey import patch_all
patch_all()

from datetime import datetime
from pymongo import MongoClient
client = MongoClient()
client.db.data.insert_many([
  {'_id': 1, 'amount': 21, 'last_updated': datetime(2020, 12, 10, 1, 3, 1), 'account': {'name': 'Customer1', 'account_number': 1}, 'txns': ['A']},
  {'_id': 2, 'amount': 16, 'last_updated': datetime(2020, 7, 23, 6, 7, 11), 'account': {'name': 'Customer2', 'account_number': 2}, 'txns': ['A', 'B']},
  {'_id': 3, 'amount': 3,  'last_updated': datetime(2021, 3, 10, 18, 43, 9), 'account': {'name': 'Customer3', 'account_number': 3}, 'txns': ['A', 'B', 'C']},
  {'_id': 4, 'amount': 0,  'last_updated': datetime(2021, 2, 25, 3, 50, 31), 'account': {'name': 'Customer4', 'account_number': 4}, 'txns': ['A', 'B', 'C', 'D']}])

Definición del esquema

PyMongoArrow se basa en un esquema de datos para ordenar los conjuntos de resultados de las consultas en formato tabular. Si no proporciona este esquema, PyMongoArrow lo infiere a partir de los datos. Puede definir el esquema creando un objeto Schema y asignando los nombres de campo a especificadores de tipo, como se muestra en el siguiente ejemplo:

from pymongoarrow.api import Schema
schema = Schema({'_id': int, 'amount': float, 'last_updated': datetime})

MongoDB utiliza documentos incrustados para representar datos anidados. PyMongoArrow ofrece compatibilidad de primera clase con estos documentos:

schema = Schema({'_id': int, 'amount': float, 'account': { 'name': str, 'account_number': int}})

PyMongoArrow también admite listas y listas anidadas:

from pyarrow import list_, string
schema = Schema({'txns': list_(string())})
polars_df = client.db.data.find_polars_all({'amount': {'$gt': 0}}, schema=schema)

Operaciones de búsqueda

El siguiente ejemplo de código muestra cómo cargar todos los registros que tienen un valor distinto de cero para el campo amount como un objeto pandas.DataFrame:

df = client.db.data.find_pandas_all({'amount': {'$gt': 0}}, schema=schema)

También puede cargar el mismo conjunto de resultados como una instancia pyarrow.Table:

arrow_table = client.db.data.find_arrow_all({'amount': {'$gt': 0}}, schema=schema)

O como ejemplo polars.DataFrame:

df = client.db.data.find_polars_all({'amount': {'$gt': 0}}, schema=schema)

O como un objeto NumPy arrays:

ndarrays = client.db.data.find_numpy_all({'amount': {'$gt': 0}}, schema=schema)

Al utilizar NumPy, el valor de retorno es un diccionario donde las claves son nombres de campo y los valores son las instancias numpy.ndarray correspondientes.

arrow_table = client.db.data.find_arrow_all({'amount': {'$gt': 0}})

Operaciones agregadas

Ejecutar una operación de agregación es similar a ejecutar una operación de búsqueda, pero se necesita una secuencia de operaciones para realizarla.

El siguiente es un ejemplo simple del método aggregate_pandas_all() que genera un nuevo marco de datos en el que todos los valores _id se agrupan y se suman sus valores amount:

df = client.db.data.aggregate_pandas_all([{'$group': {'_id': None, 'total_amount': { '$sum': '$amount' }}}])

También puede ejecutar operaciones de agregación en documentos incrustados. El siguiente ejemplo desenrolla los valores del campo anidado txn, cuenta el número de cada valor y devuelve los resultados como una lista de objetos NumPy ndarray, ordenados descendentemente:

pipeline = [{'$unwind': '$txns'}, {'$group': {'_id': '$txns', 'count': {'$sum': 1}}}, {'$sort': {"count": -1}}]
ndarrays = client.db.data.aggregate_numpy_all(pipeline)

Escribir en MongoDB

Puede utilizar el método write() para escribir objetos de los siguientes tipos en MongoDB:

• Flecha Table

• Pandas DataFrame

• NumPy ndarray

• Polars DataFrame

from pymongoarrow.api import write
from pymongo import MongoClient
coll = MongoClient().db.my_collection
write(coll, df)
write(coll, arrow_table)
write(coll, ndarrays)

data_a = [1, 2, 3]
data_b = [1, None, 3]
data = Table.from_pydict(
   {
      "a": data_a,
      "b": data_b,
   },
)
coll.drop()
write(coll, data)
col_data = list(coll.find({}))
coll.drop()
write(coll, data, exclude_none=True)
col_data_exclude_none = list(coll.find({}))
print(col_data)
print(col_data_exclude_none)

{'_id': ObjectId('...'), 'a': 1, 'b': 1}
{'_id': ObjectId('...'), 'a': 2, 'b': None}
{'_id': ObjectId('...'), 'a': 3, 'b': 3}
{'_id': ObjectId('...'), 'a': 1, 'b': 1}
{'_id': ObjectId('...'), 'a': 2}
{'_id': ObjectId('...'), 'a': 3, 'b': 3}

Escribir en otros formatos

Una vez cargados los conjuntos de resultados, puedes escribirlos en cualquier formato que admita el paquete.

Por ejemplo, para escribir la tabla referenciada por la variable arrow_table en un archivo Parquet llamado example.parquet, ejecute el siguiente código:

import pyarrow.parquet as pq
pq.write_table(arrow_table, 'example.parquet')

Pandas también admite la escritura de instancias DataFrame en una variedad de formatos, incluidos CSV y HDF. Para guardar el marco de datos referenciado por la variable df en un archivo CSV llamado out.csv, ejecute el siguiente código:

df.to_csv('out.csv', index=False)

La API de Polars es una mezcla de los dos ejemplos anteriores:

import polars as pl
df = pl.DataFrame({"foo": [1, 2, 3, 4, 5]})
df.write_parquet('example.parquet')