api.py

# Copyright 2021-present MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import multiprocessing
import warnings
from concurrent.futures import ThreadPoolExecutor
from decimal import Decimal
from typing import Literal

import numpy as np

try:
    import pandas as pd
except ImportError:
    pd = None

try:
    import polars as pl
except ImportError:
    pl = None

import pyarrow as pa
import pymongo.errors
from bson import encode
from bson.codec_options import TypeEncoder, TypeRegistry
from bson.decimal128 import Decimal128
from bson.raw_bson import RawBSONDocument
from numpy import ndarray
from pyarrow import Schema as ArrowSchema
from pyarrow import Table, timestamp
from pyarrow.types import (
    is_date32,
    is_date64,
    is_duration,
    is_float16,
    is_float32,
    is_int8,
    is_int16,
    is_list,
    is_uint8,
    is_uint16,
    is_uint32,
    is_uint64,
)
from pymongo.collection import Collection
from pymongo.common import MAX_WRITE_BATCH_SIZE
from pymongo.driver_info import DriverInfo

import pymongoarrow.version as pymongoarrow_version
from pymongoarrow.context import PyMongoArrowContext
from pymongoarrow.errors import ArrowWriteError
from pymongoarrow.result import ArrowWriteResult
from pymongoarrow.schema import Schema
from pymongoarrow.types import _validate_schema, get_numpy_type

__all__ = [
    "Schema",
    "aggregate_arrow_all",
    "aggregate_numpy_all",
    "aggregate_pandas_all",
    "aggregate_polars_all",
    "find_arrow_all",
    "find_numpy_all",
    "find_pandas_all",
    "find_polars_all",
    "write",
]


_PATCH_METHODS = [
    "aggregate_arrow_all",
    "find_arrow_all",
    "aggregate_pandas_all",
    "find_pandas_all",
    "aggregate_numpy_all",
    "find_numpy_all",
    "aggregate_polars_all",
    "find_polars_all",
]

# MongoDB 3.6's maxMessageSizeBytes minus some overhead to account
# for the command plus OP_MSG.
_MAX_MESSAGE_SIZE = 48000000 - 16 * 1024
# The maximum number of bulk write operations in one batch.
_MAX_WRITE_BATCH_SIZE = max(100000, MAX_WRITE_BATCH_SIZE)


def _add_driver_metadata(collection: Collection):
    client = collection.database.client

    if callable(client.append_metadata):
        client.append_metadata(
            DriverInfo(name="PyMongoArrow", version=pymongoarrow_version.__version__)
        )


def _process_batch(schema, codec_options, allow_invalid, batch):
    context = PyMongoArrowContext(schema, codec_options=codec_options, allow_invalid=allow_invalid)
    context.process_bson_stream(batch)

    return context.finish()


Parallelism = Literal["threads", "processes", "off"]


def find_arrow_all(
    collection,
    query,
    *,
    schema=None,
    allow_invalid=False,
    parallelism: Parallelism = "off",
    **kwargs,
):
    """Method that returns the results of a find query as a
    :class:`pyarrow.Table` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `query`: A mapping containing the query to use for the find operation.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.
      - `parallelism` (optional): Controls how batch processing is parallelized.
        Possible values are:
            - "off": (default) Disable parallelism and use the single-process behavior.
            - "threads": Always use a threaded implementation.
            - "processes": Always use a multiprocess implementation.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``find`` operation.

    :Returns:
      An instance of class:`pyarrow.Table`.
    """
    _add_driver_metadata(collection)

    for opt in ("cursor_type",):
        if kwargs.pop(opt, None):
            warnings.warn(
                f"Ignoring option {opt!r} as it is not supported by PyMongoArrow",
                UserWarning,
                stacklevel=2,
            )

    if schema:
        kwargs.setdefault("projection", schema._get_projection())

    raw_batch_cursor = collection.find_raw_batches(query, **kwargs)

    def args_iterable():
        for batch in collection.find_raw_batches(query, **kwargs):
            yield (schema, collection.codec_options, allow_invalid, batch)

    if parallelism == "threads":
        with ThreadPoolExecutor(max_workers=4) as executor:
            results = list(executor.map(lambda args: _process_batch(*args), args_iterable()))
        return pa.concat_tables(results, promote_options="permissive")

    if parallelism == "processes":
        with multiprocessing.Pool(processes=4) as pool:
            results = pool.starmap(_process_batch, args_iterable())
        return pa.concat_tables(results, promote_options="permissive")

    context = PyMongoArrowContext(
        schema, codec_options=collection.codec_options, allow_invalid=allow_invalid
    )
    for batch in raw_batch_cursor:
        context.process_bson_stream(batch)

    return context.finish()


def aggregate_arrow_all(collection, pipeline, *, schema=None, allow_invalid=False, **kwargs):
    """Method that returns the results of an aggregation pipeline as a
    :class:`pyarrow.Table` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``aggregate`` operation.
      - `pipeline`: A list of aggregation pipeline stages.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``aggregate`` operation.

    :Returns:
      An instance of class:`pyarrow.Table`.
    """
    _add_driver_metadata(collection)
    context = PyMongoArrowContext(
        schema, codec_options=collection.codec_options, allow_invalid=allow_invalid
    )

    if pipeline and ("$out" in pipeline[-1] or "$merge" in pipeline[-1]):
        msg = (
            "Aggregation pipelines containing a '$out' or '$merge' stage are "
            "not supported by PyMongoArrow"
        )
        raise ValueError(msg)

    for opt in ("batchSize", "useCursor"):
        if kwargs.pop(opt, None):
            warnings.warn(
                f"Ignoring option {opt!r} as it is not supported by PyMongoArrow",
                UserWarning,
                stacklevel=2,
            )
    if schema:
        pipeline.append({"$project": schema._get_projection()})

    raw_batch_cursor = collection.aggregate_raw_batches(pipeline, **kwargs)
    for batch in raw_batch_cursor:
        context.process_bson_stream(batch)

    return context.finish()


def _arrow_to_pandas(arrow_table):
    """Helper function that converts an Arrow Table to a Pandas DataFrame
    while minimizing peak memory consumption during conversion. The memory
    buffers backing the given Arrow Table are also destroyed after conversion.

    See https://arrow.apache.org/docs/python/pandas.html#reducing-memory-use-in-table-to-pandas
    for details.
    """
    if pd is None:
        msg = "pandas is not installed. Try pip install pandas."
        raise ValueError(msg)
    return arrow_table.to_pandas(split_blocks=True, self_destruct=True)


def find_pandas_all(
    collection,
    query,
    *,
    schema=None,
    allow_invalid=False,
    parallelism: Parallelism = "off",
    **kwargs,
):
    """Method that returns the results of a find query as a
    :class:`pandas.DataFrame` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `query`: A mapping containing the query to use for the find operation.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.
      - `parallelism` (optional): Controls how batch processing is parallelized.
        Possible values are:
            - "off": (default) Disable parallelism and use the single-process behavior.
            - "threads": Always use a threaded implementation.
            - "processes": Always use a multiprocess implementation.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``find`` operation.

    :Returns:
      An instance of class:`pandas.DataFrame`.
    """
    return _arrow_to_pandas(
        find_arrow_all(
            collection,
            query,
            schema=schema,
            allow_invalid=allow_invalid,
            parallelism=parallelism,
            **kwargs,
        )
    )


def aggregate_pandas_all(collection, pipeline, *, schema=None, allow_invalid=False, **kwargs):
    """Method that returns the results of an aggregation pipeline as a
    :class:`pandas.DataFrame` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `pipeline`: A list of aggregation pipeline stages.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``aggregate`` operation.

    :Returns:
      An instance of class:`pandas.DataFrame`.
    """
    return _arrow_to_pandas(
        aggregate_arrow_all(
            collection, pipeline, schema=schema, allow_invalid=allow_invalid, **kwargs
        )
    )


def _arrow_to_numpy(arrow_table, schema=None):
    """Helper function that converts an Arrow Table to a dictionary
    containing NumPy arrays. The memory buffers backing the given Arrow Table
    may be destroyed after conversion if the resulting Numpy array(s) is not a
    view on the Arrow data.

    See https://arrow.apache.org/docs/python/numpy.html for details.
    """
    container = {}
    schema = {i.name: i.type for i in arrow_table.schema} if not schema else schema.typemap

    for fname in schema:
        dtype = get_numpy_type(schema[fname])
        container[fname] = arrow_table[fname].to_numpy()
        if dtype == np.str_:
            container[fname] = container[fname].astype(np.str_)

    return container


def find_numpy_all(
    collection,
    query,
    *,
    schema=None,
    allow_invalid=False,
    parallelism: Parallelism = "off",
    **kwargs,
):
    """Method that returns the results of a find query as a
    :class:`dict` instance whose keys are field names and values are
    :class:`~numpy.ndarray` instances bearing the appropriate dtype.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `query`: A mapping containing the query to use for the find operation.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.
      - `parallelism` (optional): Controls how batch processing is parallelized.
         Possible values are:
            - "off": (default) Disable parallelism and use the single-process behavior.
            - "threads": Always use a threaded implementation.
            - "processes": Always use a multiprocess implementation.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``find`` operation.

    This method attempts to create each NumPy array as a view on the Arrow
    data corresponding to each field in the result set. When this is not
    possible, the underlying data is copied into a new NumPy array. See
    :meth:`pyarrow.Array.to_numpy` for more information.

    NumPy arrays returned by this method that are views on Arrow data
    are not writable. Users seeking to modify such arrays must first
    create an editable copy using :meth:`numpy.copy`.

    :Returns:
      An instance of :class:`dict`.
    """
    return _arrow_to_numpy(
        find_arrow_all(
            collection,
            query,
            schema=schema,
            allow_invalid=allow_invalid,
            parallelism=parallelism,
            **kwargs,
        ),
        schema,
    )


def aggregate_numpy_all(collection, pipeline, *, schema=None, allow_invalid=False, **kwargs):
    """Method that returns the results of an aggregation pipeline as a
    :class:`dict` instance whose keys are field names and values are
    :class:`~numpy.ndarray` instances bearing the appropriate dtype.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `query`: A mapping containing the query to use for the find operation.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``aggregate`` operation.

    This method attempts to create each NumPy array as a view on the Arrow
    data corresponding to each field in the result set. When this is not
    possible, the underlying data is copied into a new NumPy array. See
    :meth:`pyarrow.Array.to_numpy` for more information.

    NumPy arrays returned by this method that are views on Arrow data
    are not writable. Users seeking to modify such arrays must first
    create an editable copy using :meth:`numpy.copy`.

    :Returns:
      An instance of :class:`dict`.
    """
    return _arrow_to_numpy(
        aggregate_arrow_all(
            collection, pipeline, schema=schema, allow_invalid=allow_invalid, **kwargs
        ),
        schema,
    )


def _arrow_to_polars(arrow_table: pa.Table):
    """Helper function that converts an Arrow Table to a Polars DataFrame."""
    if pl is None:
        msg = "polars is not installed. Try pip install polars."
        raise ValueError(msg)

    return pl.from_arrow(arrow_table)


def find_polars_all(
    collection,
    query,
    *,
    schema=None,
    allow_invalid=False,
    parallelism: Parallelism = "off",
    **kwargs,
):
    """Method that returns the results of a find query as a
    :class:`polars.DataFrame` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `query`: A mapping containing the query to use for the find operation.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.
      - `parallelism` (optional): Controls how batch processing is parallelized.
        Possible values are:
            - "off": (default) Disable parallelism and use the single-process behavior.
            - "threads": Always use a threaded implementation.
            - "processes": Always use a multiprocess implementation.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``find`` operation.

    :Returns:
      An instance of class:`polars.DataFrame`.

    .. versionadded:: 1.3
    """
    return _arrow_to_polars(
        find_arrow_all(
            collection,
            query,
            schema=schema,
            allow_invalid=allow_invalid,
            parallelism=parallelism,
            **kwargs,
        )
    )


def aggregate_polars_all(collection, pipeline, *, schema=None, allow_invalid=False, **kwargs):
    """Method that returns the results of an aggregation pipeline as a
    :class:`polars.DataFrame` instance.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the ``find`` operation.
      - `pipeline`: A list of aggregation pipeline stages.
      - `schema` (optional): Instance of :class:`~pymongoarrow.schema.Schema`.
        If the schema is not given, it will be inferred using the data in the
        result set.
      - `allow_invalid` (optional): If set to ``True``,
        results will have all fields that do not conform to the schema silently converted to NaN.

    Additional keyword-arguments passed to this method will be passed
    directly to the underlying ``aggregate`` operation.

    :Returns:
      An instance of class:`polars.DataFrame`.
    """
    return _arrow_to_polars(
        aggregate_arrow_all(
            collection, pipeline, schema=schema, allow_invalid=allow_invalid, **kwargs
        )
    )


def _transform_bwe(bwe, offset):
    bwe["nInserted"] += offset
    for i in bwe["writeErrors"]:
        i["index"] += offset
    return {
        "writeErrors": bwe["writeErrors"],
        "nInserted": bwe["nInserted"],
        "writeConcernErrors": bwe["writeConcernErrors"],
    }


def _tabular_generator(tabular, *, exclude_none=False):
    if isinstance(tabular, Table):
        for i in tabular.to_batches():
            for row in i.to_pylist():
                if exclude_none:
                    yield {k: v for k, v in row.items() if v is not None}
                else:
                    yield row
    elif pd is not None and isinstance(tabular, pd.DataFrame):
        for row in tabular.to_dict("records"):
            if exclude_none:
                yield {k: v for k, v in row.items() if not np.isnan(v)}
            else:
                yield row
    elif pl is not None and isinstance(tabular, pl.DataFrame):
        yield from _tabular_generator(tabular.to_arrow(), exclude_none=exclude_none)
    elif isinstance(tabular, dict):
        iter_dict = {k: np.nditer(v) for k, v in tabular.items()}
        try:
            while True:
                yield {k: next(i).item() for k, i in iter_dict.items()}
        except StopIteration:
            return


class _PandasNACodec(TypeEncoder):
    """A custom type codec for Pandas NA objects."""

    @property
    def python_type(self):
        return pd.NA.__class__

    def transform_python(self, _):
        """Transform an NA object into 'None'"""
        return


class _DecimalCodec(TypeEncoder):
    """A custom type codec for Decimal objects."""

    @property
    def python_type(self):
        return Decimal

    def transform_python(self, value):
        """Transform an Decimal object into a BSON Decimal128 object"""
        return Decimal128(value)


def write(
    collection: Collection, tabular, *, exclude_none: bool = False, auto_convert: bool = True
):
    """Write data from `tabular` into the given MongoDB `collection`.

    :Parameters:
      - `collection`: Instance of :class:`~pymongo.collection.Collection`.
        against which to run the operation.
      - `tabular`: A tabular data store to use for the write operation.
      - `exclude_none`: Whether to skip writing `null` fields in documents.
      - `auto_convert` (optional): Whether to attempt a best-effort conversion of unsupported types.

    :Returns:
      An instance of :class:`result.ArrowWriteResult`.
    """
    cur_offset = 0
    results = {
        "insertedCount": 0,
    }
    tab_size = len(tabular)
    if isinstance(tabular, Table):
        # Convert date objects to datetime objects.
        changed = False
        new_types = []
        for dtype in tabular.schema.types:
            if is_date32(dtype) or is_date64(dtype):
                changed = True
                dtype = timestamp("ms")  # noqa: PLW2901
            elif auto_convert:
                if is_uint8(dtype) or is_uint16(dtype) or is_int8(dtype) or is_int16(dtype):
                    changed = True
                    dtype = pa.int32()  # noqa: PLW2901
                elif is_uint32(dtype) or is_uint64(dtype) or is_duration(dtype):
                    changed = True
                    dtype = pa.int64()  # noqa: PLW2901
                elif is_float16(dtype) or is_float32(dtype):
                    changed = True
                    dtype = pa.float64()  # noqa: PLW2901
            new_types.append(dtype)
        if changed:
            cols = [
                tabular.column(i).cast(new_types[i])
                if not is_list(new_types[i])
                else tabular.column(i)
                for i in range(tabular.num_columns)
            ]
            tabular = Table.from_arrays(cols, names=tabular.column_names)
        _validate_schema(tabular.schema.types)
    elif pd is not None and isinstance(tabular, pd.DataFrame):
        _validate_schema(ArrowSchema.from_pandas(tabular).types)
    elif pl is not None and isinstance(tabular, pl.DataFrame):
        tabular = tabular.to_arrow()  # zero-copy in most cases and done in tabular_gen anyway
        _validate_schema(tabular.schema.types)
    elif (
        isinstance(tabular, dict)
        and len(tabular.values()) >= 1
        and ndarray is not None
        and all(isinstance(i, ndarray) for i in tabular.values())
    ):
        _validate_schema([i.dtype for i in tabular.values()])
        tab_size = len(next(iter(tabular.values())))
    else:
        msg = (
            f"Invalid tabular data object of type {type(tabular)} \n"
            "Please ensure that it is one of the supported types: "
            "DataFrame, Table, or a dictionary containing NumPy arrays."
        )
        raise ValueError(msg)

    tabular_gen = _tabular_generator(tabular, exclude_none=exclude_none)

    # Add handling for special case types.
    codec_options = collection.codec_options
    base_codecs = []
    if hasattr(codec_options.type_registry, "codecs"):
        base_codecs = codec_options.type_registry.codecs
    if pd is not None:
        type_registry = TypeRegistry([*base_codecs, _PandasNACodec(), _DecimalCodec()])
    else:
        type_registry = TypeRegistry([*base_codecs, _DecimalCodec()])
    codec_options = codec_options.with_options(type_registry=type_registry)
    _add_driver_metadata(collection)

    while cur_offset < tab_size:
        cur_size = 0
        cur_batch = []
        i = 0
        while (
            cur_size <= _MAX_MESSAGE_SIZE
            and len(cur_batch) <= _MAX_WRITE_BATCH_SIZE
            and cur_offset + i < tab_size
        ):
            enc_tab = RawBSONDocument(encode(next(tabular_gen), codec_options=codec_options))
            cur_batch.append(enc_tab)
            cur_size += len(enc_tab.raw)
            i += 1
        try:
            collection.insert_many(cur_batch)
        except pymongo.errors.BulkWriteError as bwe:
            raise ArrowWriteError(_transform_bwe(dict(bwe.details), cur_offset)) from bwe
        except pymongo.errors.PyMongoError as pme:
            raise ArrowWriteError(
                {
                    "writeErrors": [{"errmsg": str(pme), "index": cur_offset}],
                    "nInserted": cur_offset,
                    "writeConcernErrors": [],
                }
            ) from pme
        results["insertedCount"] += i
        cur_offset += i

    return ArrowWriteResult(results)