iterator.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 warnings
from functools import cached_property
from itertools import islice, repeat
from typing import Iterable, Tuple

from nanoarrow._lib import CArrayView, CArrowType
from nanoarrow.c_array_stream import c_array_stream
from nanoarrow.c_schema import c_schema, c_schema_view
from nanoarrow.schema import Schema


def iter_py(obj, schema=None) -> Iterable:
    """Iterate over items in zero or more arrays

    Returns an iterator over an array stream where each item is a
    Python representation of the next element.

    Paramters
    ---------
    obj : array stream-like
        An array-like or array stream-like object as sanitized by
        :func:`c_array_stream`.
    schema : schema-like, optional
        An optional schema, passed to :func:`c_array_stream`.

    Examples
    --------

    >>> import nanoarrow as na
    >>> from nanoarrow import iterator
    >>> array = na.c_array([1, 2, 3], na.int32())
    >>> list(iterator.iter_py(array))
    [1, 2, 3]
    """
    return PyIterator.get_iterator(obj, schema=schema)


def iter_tuples(obj, schema=None) -> Iterable[Tuple]:
    """Iterate over rows in zero or more struct arrays

    Returns an iterator over an array stream of struct arrays (i.e.,
    record batches) where each item is a tuple of the items in each
    row. This is different than :func:`iter_py`, which encodes struct
    columns as dictionaries.

    Paramters
    ---------
    obj : array stream-like
        An array-like or array stream-like object as sanitized by
        :func:`c_array_stream`.
    schema : schema-like, optional
        An optional schema, passed to :func:`c_array_stream`.

    Examples
    --------

    >>> import nanoarrow as na
    >>> from nanoarrow import iterator
    >>> import pyarrow as pa
    >>> array = pa.record_batch([pa.array([1, 2, 3])], names=["col1"])
    >>> list(iterator.iter_tuples(array))
    [(1,), (2,), (3,)]
    """
    return RowTupleIterator.get_iterator(obj, schema=schema)


def iter_array_views(obj, schema=None) -> Iterable[CArrayView]:
    """Iterate over prepared views of each array

    Returns an iterator which yields a :func:`c_array_view`
    for each chunk in ``obj``.

    Paramters
    ---------
    obj : array stream-like
        An array-like or array stream-like object as sanitized by
        :func:`c_array_stream`.
    schema : schema-like, optional
        An optional schema, passed to :func:`c_array_stream`.

    Examples
    --------

    >>> import nanoarrow as na
    >>> from nanoarrow import iterator
    >>> array = na.c_array([1, 2, 3], na.int32())
    >>> list(iterator.iter_array_views(array))
    [<nanoarrow.c_array.CArrayView>
    - storage_type: 'int32'
    - length: 3
    - offset: 0
    - null_count: 0
    - buffers[2]:
      - validity <bool[0 b] >
      - data <int32[12 b] 1 2 3>
    - dictionary: NULL
    - children[0]:]
    """
    with c_array_stream(obj, schema) as stream:
        for array in stream:
            yield array.view()


class InvalidArrayWarning(UserWarning):
    pass


class LossyConversionWarning(UserWarning):
    pass


class UnregisteredExtensionWarning(UserWarning):
    pass


class ArrayViewBaseIterator:
    """Base class for iterators and visitors that use an internal ArrowArrayView
    as the basis for conversion to Python objects. Intended for internal use.
    """

    def __init__(self, schema, *, array_view=None):
        self._schema = c_schema(schema)
        self._schema_view = c_schema_view(schema)

        if array_view is None:
            self._array_view = CArrayView.from_schema(self._schema)
        else:
            self._array_view = array_view

    @cached_property
    def schema(self) -> Schema:
        return Schema(self._schema)

    @cached_property
    def _object_label(self):
        if self._schema.name:
            return f"{self._schema.name} <{self._schema_view.type}>"
        else:
            return f"<unnamed {self._schema_view.type}>"

    def _contains_nulls(self):
        return self._schema_view.nullable and self._array_view.null_count != 0

    def _set_array(self, array):
        self._array_view._set_array(array)
        return self

    def _warn(self, message, category):
        warnings.warn(f"{self._object_label}: {message}", category)


class PyIterator(ArrayViewBaseIterator):
    """Iterate over the Python object version of values in an ArrowArrayView.
    Intended for internal use.
    """

    @classmethod
    def get_iterator(cls, obj, schema=None):
        with c_array_stream(obj, schema=schema) as stream:
            iterator = cls(stream._get_cached_schema())
            for array in stream:
                iterator._set_array(array)
                yield from iterator

    def __init__(self, schema, *, array_view=None):
        super().__init__(schema, array_view=array_view)

        self._children = list(
            map(self._make_child, self._schema.children, self._array_view.children)
        )

        if self._schema.dictionary is None:
            self._dictionary = None
        else:
            self._dictionary = self._make_child(
                self._schema.dictionary, self._array_view.dictionary
            )

    def _make_child(self, schema, array_view):
        return type(self)(schema, array_view=array_view)

    @cached_property
    def _child_names(self):
        return [child.name for child in self._schema.children]

    def __iter__(self):
        """Iterate over all elements in the current chunk"""
        return self._iter_chunk(0, len(self._array_view))

    def _iter_chunk(self, offset, length):
        """Iterate over all elements in a slice of the current chunk"""
        # Check for an extension type first since this isn't reflected by
        # self._schema_view.type_id. Currently we just return the storage
        # iterator with a warning for extension types.
        maybe_extension_name = self._schema_view.extension_name
        if maybe_extension_name:
            self._warn(
                f"Converting unregistered extension '{maybe_extension_name}' "
                "as storage type",
                UnregisteredExtensionWarning,
            )

        type_id = self._schema_view.type_id
        if type_id not in _ITEMS_ITER_LOOKUP:
            raise KeyError(
                f"Can't resolve iterator for type '{self._schema_view.type}'"
            )

        factory = getattr(self, _ITEMS_ITER_LOOKUP[type_id])
        return factory(offset, length)

    def _dictionary_iter(self, offset, length):
        dictionary = list(
            self._dictionary._iter_chunk(0, len(self._dictionary._array_view))
        )
        for dict_index in self._primitive_iter(offset, length):
            yield None if dict_index is None else dictionary[dict_index]

    def _wrap_iter_nullable(self, validity, items):
        for is_valid, item in zip(validity, items):
            yield item if is_valid else None

    def _struct_tuple_iter(self, offset, length):
        view = self._array_view
        offset += view.offset
        items = zip(*(child._iter_chunk(offset, length) for child in self._children))

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            return self._wrap_iter_nullable(validity, items)
        else:
            return items

    def _struct_iter(self, offset, length):
        names = self._child_names
        for item in self._struct_tuple_iter(offset, length):
            yield None if item is None else {key: val for key, val in zip(names, item)}

    def _list_iter(self, offset, length):
        view = self._array_view
        offset += view.offset

        offsets = memoryview(view.buffer(1))[offset : (offset + length + 1)]
        starts = offsets[:-1]
        ends = offsets[1:]
        child = self._children[0]
        child_iter = child._iter_chunk(starts[0], ends[-1] - starts[0])

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            for is_valid, start, end in zip(validity, starts, ends):
                item = list(islice(child_iter, end - start))
                yield item if is_valid else None
        else:
            for start, end in zip(starts, ends):
                yield list(islice(child_iter, end - start))

    def _fixed_size_list_iter(self, offset, length):
        view = self._array_view
        offset += view.offset
        child = self._children[0]
        fixed_size = view.layout.child_size_elements
        child_iter = child._iter_chunk(offset * fixed_size, length * fixed_size)

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            for is_valid in validity:
                item = list(islice(child_iter, fixed_size))
                yield item if is_valid else None
        else:
            for _ in range(length):
                yield list(islice(child_iter, fixed_size))

    def _string_iter(self, offset, length):
        view = self._array_view
        offset += view.offset
        offsets = memoryview(view.buffer(1))[offset : (offset + length + 1)]
        starts = offsets[:-1]
        ends = offsets[1:]
        data = memoryview(view.buffer(2))

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            for is_valid, start, end in zip(validity, starts, ends):
                yield str(data[start:end], "UTF-8") if is_valid else None
        else:
            for start, end in zip(starts, ends):
                yield str(data[start:end], "UTF-8")

    def _binary_iter(self, offset, length):
        view = self._array_view
        offset += view.offset
        offsets = memoryview(view.buffer(1))[offset : (offset + length + 1)]
        starts = offsets[:-1]
        ends = offsets[1:]
        data = memoryview(view.buffer(2))

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            for is_valid, start, end in zip(validity, starts, ends):
                yield bytes(data[start:end]) if is_valid else None
        else:
            for start, end in zip(starts, ends):
                yield bytes(data[start:end])

    def _decimal_iter(self, offset, length):
        from decimal import Context, Decimal
        from sys import byteorder

        storage = self._primitive_iter(offset, length)
        precision = self._schema_view.decimal_precision

        # The approach here it to use Decimal(<integer>).scaleb(-scale),
        # which is a balance between simplicity, performance, and
        # safety (ensuring that we stay independent from the global precision).
        # We cache the scaleb and context to avoid doing so in the loop (the
        # argument to scaleb is transformed to a decimal by the .scaleb()
        # implementation).
        #
        # It would probably be fastest to go straight from binary
        # to string to decimal, since creating a decimal from a string
        # appears to be the fastest constructor.
        scaleb = Decimal(-self._schema_view.decimal_scale)
        context = Context(prec=precision)

        for item in storage:
            if item is None:
                yield None
            else:
                int_value = int.from_bytes(item, byteorder)
                yield Decimal(int_value).scaleb(scaleb, context)

    def _date_iter(self, offset, length):
        from datetime import date, timedelta

        storage = self._primitive_iter(offset, length)
        epoch = date(1970, 1, 1)

        if self._schema_view.type_id == CArrowType.DATE32:
            for item in storage:
                if item is None:
                    yield item
                else:
                    yield epoch + timedelta(item)
        else:
            for item in storage:
                if item is None:
                    yield item
                else:
                    yield epoch + timedelta(milliseconds=item)

    def _time_iter(self, offset, length):
        from datetime import time

        for item in self._iter_time_components(offset, length):
            if item is None:
                yield None
            else:
                days, hours, mins, secs, us = item
                if days != 0:
                    self._warn("days != 0", InvalidArrayWarning)

                yield time(hours, mins, secs, us)

    def _timestamp_iter(self, offset, length):
        from datetime import datetime

        epoch = datetime(1970, 1, 1, tzinfo=_get_tzinfo("UTC"))
        parent = self._duration_iter(offset, length)

        tz = self._schema_view.timezone
        if tz:
            tz = _get_tzinfo(tz)

            for item in parent:
                if item is None:
                    yield None
                else:
                    yield (epoch + item).astimezone(tz)
        else:
            epoch = epoch.replace(tzinfo=None)
            for item in parent:
                if item is None:
                    yield None
                else:
                    yield epoch + item

    def _duration_iter(self, offset, length):
        from datetime import timedelta

        storage = self._primitive_iter(offset, length)

        unit = self._schema_view.time_unit
        if unit == "s":
            to_us = 1_000_000
        elif unit == "ms":
            to_us = 1000
        elif unit == "us":
            to_us = 1
        elif unit == "ns":
            storage = self._iter_us_from_ns(storage)
            to_us = 1

        for item in storage:
            if item is None:
                yield None
            else:
                yield timedelta(microseconds=item * to_us)

    def _iter_time_components(self, offset, length):
        storage = self._primitive_iter(offset, length)

        unit = self._schema_view.time_unit
        if unit == "s":
            to_us = 1_000_000
        elif unit == "ms":
            to_us = 1000
        elif unit == "us":
            to_us = 1
        elif unit == "ns":
            storage = self._iter_us_from_ns(storage)
            to_us = 1

        us_per_sec = 1_000_000
        us_per_min = us_per_sec * 60
        us_per_hour = us_per_min * 60
        us_per_day = us_per_hour * 24

        for item in storage:
            if item is None:
                yield None
            else:
                us = item * to_us
                days = us // us_per_day
                us = us % us_per_day
                hours = us // us_per_hour
                us = us % us_per_hour
                mins = us // us_per_min
                us = us % us_per_min
                secs = us // us_per_sec
                us = us % us_per_sec
                yield days, hours, mins, secs, us

    def _iter_us_from_ns(self, parent):
        for item in parent:
            if item is None:
                yield None
            else:
                if item % 1000 != 0:
                    self._warn("nanoseconds discarded", LossyConversionWarning)
                yield item // 1000

    def _primitive_iter(self, offset, length):
        view = self._array_view
        offset += view.offset
        items = view.buffer(1).elements(offset, length)

        if self._contains_nulls():
            validity = view.buffer(0).elements(offset, length)
            return self._wrap_iter_nullable(validity, items)
        else:
            return iter(items)

    def _null_iter(self, offset, length):
        return repeat(None, length)


class RowTupleIterator(PyIterator):
    """Iterate over rows of a struct array (stream) where each row is a
    tuple instead of a dictionary. This is usually faster and matches other
    Python concepts more closely (e.g., dbapi's cursor, pandas itertuples).
    Intended for internal use.
    """

    def __init__(self, schema, *, array_view=None):
        super().__init__(schema, array_view=array_view)
        if self._schema_view.type != "struct":
            raise TypeError(
                "RowTupleIterator can only iterate over struct arrays "
                f"(got '{self._schema_view.type}')"
            )

    def _make_child(self, schema, array_view):
        return PyIterator(schema, array_view=array_view)

    def _iter_chunk(self, offset, length):
        return self._struct_tuple_iter(offset, length)


def _get_tzinfo(tz_string, strategy=None):
    import re
    from datetime import timedelta, timezone

    # We can handle UTC without any imports
    if tz_string.upper() == "UTC":
        return timezone.utc

    # Arrow also allows fixed-offset in the from +HH:MM
    maybe_fixed_offset = re.search(r"^([+-])([0-9]{2}):([0-9]{2})$", tz_string)
    if maybe_fixed_offset:
        sign, hours, minutes = maybe_fixed_offset.groups()
        sign = 1 if sign == "+" else -1
        return timezone(sign * timedelta(hours=int(hours), minutes=int(minutes)))

    # Try zoneinfo.ZoneInfo() (Python 3.9+)
    if strategy is None or "zoneinfo" in strategy:
        try:
            from zoneinfo import ZoneInfo

            return ZoneInfo(tz_string)
        except ImportError:
            pass

    # Try dateutil.tz.gettz()
    if strategy is None or "dateutil" in strategy:
        try:
            from dateutil.tz import gettz

            return gettz(tz_string)
        except ImportError:
            pass

    raise RuntimeError(
        "zoneinfo (Python 3.9+, with tzdata on Windows) or "
        "dateutil is required to resolve timezone"
    )


_ITEMS_ITER_LOOKUP = {
    CArrowType.NA: "_null_iter",
    CArrowType.BINARY: "_binary_iter",
    CArrowType.LARGE_BINARY: "_binary_iter",
    CArrowType.STRING: "_string_iter",
    CArrowType.LARGE_STRING: "_string_iter",
    CArrowType.STRUCT: "_struct_iter",
    CArrowType.LIST: "_list_iter",
    CArrowType.LARGE_LIST: "_list_iter",
    CArrowType.FIXED_SIZE_LIST: "_fixed_size_list_iter",
    CArrowType.DICTIONARY: "_dictionary_iter",
    CArrowType.DATE32: "_date_iter",
    CArrowType.DATE64: "_date_iter",
    CArrowType.TIME32: "_time_iter",
    CArrowType.TIME64: "_time_iter",
    CArrowType.TIMESTAMP: "_timestamp_iter",
    CArrowType.DURATION: "_duration_iter",
    CArrowType.DECIMAL128: "_decimal_iter",
    CArrowType.DECIMAL256: "_decimal_iter",
}

_PRIMITIVE_TYPE_NAMES = [
    "BOOL",
    "UINT8",
    "INT8",
    "UINT16",
    "INT16",
    "UINT32",
    "INT32",
    "UINT64",
    "INT64",
    "HALF_FLOAT",
    "FLOAT",
    "DOUBLE",
    "FIXED_SIZE_BINARY",
    "INTERVAL_MONTHS",
    "INTERVAL_DAY_TIME",
    "INTERVAL_MONTH_DAY_NANO",
]

for type_name in _PRIMITIVE_TYPE_NAMES:
    type_id = getattr(CArrowType, type_name)
    _ITEMS_ITER_LOOKUP[type_id] = "_primitive_iter"