odps/tunnel/io/reader.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2025 Alibaba Group Holding Ltd.
#
# 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 collections
import functools
import json
import struct
import sys
import warnings
from collections import OrderedDict
from decimal import Decimal
from io import BytesIO, IOBase, StringIO

try:
    import numpy as np
except ImportError:
    np = None
try:
    import pandas as pd
except (ImportError, ValueError):
    pd = None
try:
    import pyarrow as pa
except (AttributeError, ImportError):
    pa = None
try:
    import pyarrow.compute as pac
except (AttributeError, ImportError):
    pac = None

from ... import compat, types, utils
from ...config import options
from ...errors import ChecksumError, DatetimeOverflowError
from ...lib.monotonic import monotonic
from ...models import Record
from ...readers import AbstractRecordReader
from ...types import PartitionSpec
from ..base import TunnelMetrics
from ..checksum import Checksum
from ..pb import wire_format
from ..pb.decoder import Decoder
from ..pb.errors import DecodeError
from ..wireconstants import ProtoWireConstants
from .types import odps_schema_to_arrow_schema

try:
    if not options.force_py:
        from .reader_c import BaseTunnelRecordReader, convert_legacy_decimal_bytes
    else:
        BaseTunnelRecordReader = convert_legacy_decimal_bytes = None
except ImportError as e:
    if options.force_c:
        raise e
    BaseTunnelRecordReader = convert_legacy_decimal_bytes = None

MICRO_SEC_PER_SEC = 1000000

if BaseTunnelRecordReader is None:

    class BaseTunnelRecordReader(AbstractRecordReader):
        def __init__(
            self,
            schema,
            stream_creator,
            columns=None,
            partition_spec=None,
            append_partitions=False,
        ):
            self._schema = schema
            if columns is None:
                self._columns = (
                    self._schema.columns
                    if append_partitions
                    else self._schema.simple_columns
                )
            else:
                self._columns = [self._schema[c] for c in columns]

            self._enable_client_metrics = options.tunnel.enable_client_metrics
            self._server_metrics_string = None
            self._local_wall_time_ms = 0
            self._acc_network_time_ms = 0
            self._injected_error = None

            self._curr_cursor = 0
            self._stream_creator = stream_creator
            self._reader = None
            self._reopen_reader()

            if self._enable_client_metrics:
                ts = monotonic()

            self._read_limit = options.table_read_limit
            self._to_datetime = utils.MillisecondsConverter().from_milliseconds
            self._to_datetime_utc = utils.MillisecondsConverter(
                local_tz=False
            ).from_milliseconds
            self._to_date = utils.to_date
            self._partition_spec = (
                PartitionSpec(partition_spec) if partition_spec else None
            )
            self._append_partitions = append_partitions

            map_as_ordered_dict = options.map_as_ordered_dict
            if map_as_ordered_dict is None:
                map_as_ordered_dict = sys.version_info[:2] <= (3, 6)
            self._map_dict_hook = OrderedDict if map_as_ordered_dict else dict

            struct_as_ordered_dict = options.struct_as_ordered_dict
            if struct_as_ordered_dict is None:
                struct_as_ordered_dict = sys.version_info[:2] <= (3, 6)
            self._struct_dict_hook = OrderedDict if struct_as_ordered_dict else dict

            if self._enable_client_metrics:
                self._local_wall_time_ms += compat.long_type(
                    MICRO_SEC_PER_SEC * (monotonic() - ts)
                )

        def _mode(self):
            return "py"

        @property
        def count(self):
            return self._curr_cursor

        @property
        def _network_wall_time_ms(self):
            return self._reader.network_wall_time_ms + self._acc_network_time_ms

        def _inject_error(self, cursor, exc):
            self._injected_error = (cursor, exc)

        def _reopen_reader(self):
            if self._enable_client_metrics:
                ts = monotonic()

            stream = self._stream_creator(self._curr_cursor)
            if self._enable_client_metrics:
                self._acc_network_time_ms += compat.long_type(
                    MICRO_SEC_PER_SEC * (monotonic() - ts)
                )
                if self._reader is not None:
                    self._acc_network_time_ms += self._reader.network_wall_time_ms

            self._last_n_bytes = len(self._reader) if self._curr_cursor != 0 else 0
            self._reader = Decoder(
                stream, record_network_time=self._enable_client_metrics
            )
            self._crc = Checksum()
            self._crccrc = Checksum()
            self._attempt_row_count = 0

            if self._enable_client_metrics:
                self._local_wall_time_ms += compat.long_type(
                    MICRO_SEC_PER_SEC * (monotonic() - ts)
                )

        def _read_field(self, data_type):
            if data_type == types.float_:
                val = self._reader.read_float()
                self._crc.update_float(val)
            elif data_type == types.double:
                val = self._reader.read_double()
                self._crc.update_double(val)
            elif data_type == types.boolean:
                val = self._reader.read_bool()
                self._crc.update_bool(val)
            elif data_type in types.integer_types:
                val = self._reader.read_sint64()
                self._crc.update_long(val)
            elif data_type == types.string:
                val = self._reader.read_string()
                self._crc.update(val)
            elif data_type == types.binary:
                val = self._reader.read_string()
                self._crc.update(val)
            elif data_type == types.datetime:
                val = self._reader.read_sint64()
                self._crc.update_long(val)
                try:
                    val = self._to_datetime(val)
                except DatetimeOverflowError:
                    if not options.tunnel.overflow_date_as_none:
                        raise
                    val = None
            elif data_type == types.date:
                val = self._reader.read_sint64()
                self._crc.update_long(val)
                val = self._to_date(val)
            elif data_type == types.timestamp or data_type == types.timestamp_ntz:
                to_datetime = (
                    self._to_datetime_utc
                    if data_type == types.timestamp_ntz
                    else self._to_datetime
                )
                l_val = self._reader.read_sint64()
                self._crc.update_long(l_val)
                nano_secs = self._reader.read_sint32()
                self._crc.update_int(nano_secs)
                if pd is None:
                    raise ImportError(
                        "To use TIMESTAMP in pyodps, you need to install pandas."
                    )
                try:
                    val = pd.Timestamp(to_datetime(l_val * 1000)) + pd.Timedelta(
                        nanoseconds=nano_secs
                    )
                except DatetimeOverflowError:
                    if not options.tunnel.overflow_date_as_none:
                        raise
                    val = None
            elif data_type == types.interval_day_time:
                l_val = self._reader.read_sint64()
                self._crc.update_long(l_val)
                nano_secs = self._reader.read_sint32()
                self._crc.update_int(nano_secs)
                if pd is None:
                    raise ImportError(
                        "To use INTERVAL_DAY_TIME in pyodps, you need to install pandas."
                    )
                val = pd.Timedelta(seconds=l_val, nanoseconds=nano_secs)
            elif data_type == types.interval_year_month:
                l_val = self._reader.read_sint64()
                self._crc.update_long(l_val)
                return compat.Monthdelta(l_val)
            elif data_type == types.json:
                sval = self._reader.read_string()
                val = json.loads(sval)
                self._crc.update(sval)
            elif isinstance(data_type, (types.Char, types.Varchar)):
                val = self._reader.read_string()
                self._crc.update(val)
            elif isinstance(data_type, types.Decimal):
                val = self._reader.read_string()
                self._crc.update(val)
            elif isinstance(data_type, types.Array):
                val = self._read_array(data_type.value_type)
            elif isinstance(data_type, types.Map):
                keys = self._read_array(data_type.key_type)
                values = self._read_array(data_type.value_type)
                val = self._map_dict_hook(zip(keys, values))
            elif isinstance(data_type, types.Struct):
                val = self._read_struct(data_type)
            else:
                raise IOError("Unsupported type %s" % data_type)
            return val

        def _read_array(self, value_type):
            res = []

            size = self._reader.read_uint32()
            for _ in range(size):
                if self._reader.read_bool():
                    res.append(None)
                else:
                    res.append(self._read_field(value_type))

            return res

        def _read_struct(self, value_type):
            res_list = [None] * len(value_type.field_types)
            for idx, field_type in enumerate(value_type.field_types.values()):
                if not self._reader.read_bool():
                    res_list[idx] = self._read_field(field_type)
            if options.struct_as_dict:
                return self._struct_dict_hook(
                    zip(value_type.field_types.keys(), res_list)
                )
            else:
                return value_type.namedtuple_type(*res_list)

        def _read_single_record(self):
            if (
                self._injected_error is not None
                and self._injected_error[0] == self._curr_cursor
            ):
                self._injected_error = None
                raise self._injected_error[1]

            if self._read_limit is not None and self.count >= self._read_limit:
                warnings.warn(
                    "Number of lines read via tunnel already reaches the limitation.",
                    RuntimeWarning,
                )
                return None

            record = Record(self._columns, max_field_size=(1 << 63) - 1)

            while True:
                index, _ = self._reader.read_field_number_and_wire_type()

                if index == 0:
                    continue
                if index == ProtoWireConstants.TUNNEL_END_RECORD:
                    checksum = utils.long_to_int(self._crc.getvalue())
                    if int(self._reader.read_uint32()) != utils.int_to_uint(checksum):
                        raise ChecksumError("Checksum invalid")
                    self._crc.reset()
                    self._crccrc.update_int(checksum)
                    break

                if index == ProtoWireConstants.TUNNEL_META_COUNT:
                    if self._attempt_row_count != self._reader.read_sint64():
                        raise IOError("count does not match")

                    (
                        idx_of_checksum,
                        wire_type,
                    ) = self._reader.read_field_number_and_wire_type()
                    if ProtoWireConstants.TUNNEL_META_CHECKSUM != idx_of_checksum:
                        if wire_type != wire_format.WIRETYPE_LENGTH_DELIMITED:
                            raise IOError("Invalid stream data.")
                        self._crc.update_int(idx_of_checksum)

                        self._server_metrics_string = self._reader.read_string()
                        self._crc.update(self._server_metrics_string)

                        idx_of_checksum = (
                            self._reader.read_field_number_and_wire_type()[0]
                        )
                        if idx_of_checksum != ProtoWireConstants.TUNNEL_END_METRICS:
                            raise IOError("Invalid stream data.")
                        checksum = utils.long_to_int(self._crc.getvalue())
                        self._crc.reset()
                        if utils.int_to_uint(checksum) != int(
                            self._reader.read_uint32()
                        ):
                            raise ChecksumError("Checksum invalid.")

                        idx_of_checksum = (
                            self._reader.read_field_number_and_wire_type()[0]
                        )
                    if ProtoWireConstants.TUNNEL_META_CHECKSUM != idx_of_checksum:
                        raise IOError("Invalid stream data.")
                    if int(self._crccrc.getvalue()) != self._reader.read_uint32():
                        raise ChecksumError("Checksum invalid.")

                    return

                if index > len(self._columns):
                    raise IOError(
                        "Invalid protobuf tag. Perhaps the datastream "
                        "from server is crushed."
                    )

                self._crc.update_int(index)

                i = index - 1
                record[i] = self._read_field(self._columns[i].type)

            if self._append_partitions and self._partition_spec is not None:
                for k, v in self._partition_spec.items():
                    try:
                        record[k] = v
                    except KeyError:
                        # skip non-existing fields
                        pass

            self._curr_cursor += 1
            self._attempt_row_count += 1
            return record

        def read(self):
            if self._enable_client_metrics:
                ts = monotonic()
            result = utils.call_with_retry(
                self._read_single_record, reset_func=self._reopen_reader
            )
            if self._enable_client_metrics:
                self._local_wall_time_ms += int(MICRO_SEC_PER_SEC * (monotonic() - ts))
            return result

        def reads(self):
            return self.__iter__()

        @property
        def n_bytes(self):
            return self._last_n_bytes + len(self._reader)

        def get_total_bytes(self):
            return self.n_bytes


class TunnelRecordReader(BaseTunnelRecordReader, AbstractRecordReader):
    """
    Reader object to read data from ODPS in records. Should be created
    with :meth:`TableDownloadSession.open_record_reader`.

    :Example:

    .. code-block:: python

        from odps.tunnel import TableTunnel

            tunnel = TableTunnel(o)
            download_session = tunnel.create_download_session('my_table', partition_spec='pt=test')

            # create a TunnelRecordReader
            with download_session.open_record_reader(0, download_session.count) as reader:
                for record in reader:
                    print(record.values)
    """

    def __next__(self):
        record = self.read()
        if record is None:
            raise StopIteration
        return record

    next = __next__

    @property
    def schema(self):
        return self._schema

    @property
    def metrics(self):
        if self._server_metrics_string is None:
            return None
        return TunnelMetrics.from_server_json(
            type(self).__name__,
            self._server_metrics_string,
            self._local_wall_time_ms,
            self._network_wall_time_ms,
        )

    def close(self):
        if hasattr(self._schema, "close"):
            self._schema.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


try:
    TunnelRecordReader.read.__doc__ = """
    Read next record.

    :return: A record object
    :rtype: :class:`~odps.models.Record`
    """
except:
    pass


class ArrowStreamReader(IOBase):
    def __init__(self, raw_reader, arrow_schema):
        self._reader = raw_reader

        self._crc = Checksum()
        self._crccrc = Checksum()
        self._pos = 0
        self._chunk_size = None

        self._buffers = collections.deque()
        self._buffers.append(BytesIO(arrow_schema.serialize().to_pybytes()))

    def __len__(self):
        return len(self._reader)

    def readable(self):
        return True

    @staticmethod
    def _read_unint32(b):
        return struct.unpack("!I", b)

    def _read_chunk_size(self):
        try:
            i = self._read_unint32(self._reader.read(4))
            self._pos += 4
            return i[0]  # unpack() result is a 1-element tuple.
        except struct.error as e:
            raise DecodeError(e)

    def _read_chunk(self):
        read_size = self._chunk_size + 4
        b = self._reader.read(read_size)
        if 0 < len(b) < 4:
            raise ChecksumError("Checksum invalid")
        self._pos += len(b)
        self._crc.update(b[:-4])
        self._crccrc.update(b[:-4])
        return b

    def _fill_next_buffer(self):
        if self._chunk_size is None:
            self._chunk_size = self._read_chunk_size()

        b = self._read_chunk()
        data = b[:-4]
        crc_data = b[-4:]
        if len(b) == 0:
            return

        if len(b) < self._chunk_size + 4:
            # is last chunk
            read_checksum = self._read_unint32(crc_data)[0]
            checksum = int(self._crccrc.getvalue())
            if checksum != read_checksum:
                raise ChecksumError("Checksum invalid")
            self._pos += len(data) + 4
            self._buffers.append(BytesIO(data))
            self._crccrc.reset()
        else:
            checksum = int(self._crc.getvalue())
            read_checksum = self._read_unint32(crc_data)[0]
            if checksum != read_checksum:
                raise ChecksumError("Checksum invalid")
            self._crc.reset()
            self._buffers.append(BytesIO(data))

    def read(self, nbytes=None):
        tot_size = 0
        bufs = []
        while nbytes is None or tot_size < nbytes:
            if not self._buffers:
                self._fill_next_buffer()
                if not self._buffers:
                    break

            to_read = nbytes - tot_size if nbytes is not None else None
            buf = self._buffers[0].read(to_read)
            if not buf:
                self._buffers.popleft()
            else:
                bufs.append(buf)
                tot_size += len(buf)
        if len(bufs) == 1:
            return bufs[0]
        return b"".join(bufs)

    def close(self):
        if hasattr(self._reader, "close"):
            self._reader.close()


class TunnelArrowReader(object):
    """
    Reader object to read data from ODPS in Arrow format. Should be created
    with :meth:`TableDownloadSession.open_arrow_reader`.

    :Example:

    .. code-block:: python

        from odps.tunnel import TableTunnel

            tunnel = TableTunnel(o)
            download_session = tunnel.create_download_session('my_table', partition_spec='pt=test')

            # create a TunnelArrowReader
            with download_session.open_arrow_reader(0, download_session.count) as reader:
                for batch in reader:
                    print(batch.to_pandas())
    """

    def __init__(
        self,
        schema,
        stream_creator,
        columns=None,
        partition_spec=None,
        append_partitions=False,
        use_ipc_stream=False,
    ):
        if pa is None:
            raise ValueError("To use arrow reader you need to install pyarrow")

        self._raw_schema = schema

        raw_arrow_schema = odps_schema_to_arrow_schema(schema)
        if columns is None:
            self._schema = schema
            self._arrow_schema = self._raw_arrow_schema = raw_arrow_schema
        else:
            self._schema = types.OdpsSchema([schema[c] for c in columns])
            self._raw_arrow_schema = pa.schema(
                [s for s in raw_arrow_schema if s.name in columns]
            )
            self._arrow_schema = odps_schema_to_arrow_schema(self._schema)
        self._columns = columns

        self._append_partitions = append_partitions
        self._partition_spec = partition_spec

        self._pos = 0
        self._stream_creator = stream_creator
        self._use_ipc_stream = use_ipc_stream
        self._reopen_reader()

        self._to_datetime = utils.MillisecondsConverter().from_milliseconds
        self._read_limit = options.table_read_limit

        self.closed = False

        self._pd_column_converters = dict()
        # True if need to convert numeric columns with None as float types
        self._coerce_numpy_columns = set()
        for col in schema.simple_columns:
            arrow_type = self._arrow_schema.field(
                self._arrow_schema.get_field_index(col.name)
            ).type
            if isinstance(col.type, (types.Map, types.Array, types.Struct)):
                self._pd_column_converters[col.name] = ArrowRecordFieldConverter(
                    col.type, convert_ts=False
                )
            elif options.tunnel.pd_cast_mode == "numpy" and (
                pa.types.is_integer(arrow_type) or pa.types.is_floating(arrow_type)
            ):
                self._coerce_numpy_columns.add(col.name)

        self._injected_error = None

    def _reopen_reader(self):
        self._last_n_bytes = len(self._reader) if self._pos != 0 else 0
        input_stream = self._stream_creator(self._pos)
        self._arrow_stream = None
        if self._use_ipc_stream:
            self._reader = input_stream
        else:
            self._reader = ArrowStreamReader(input_stream, self._raw_arrow_schema)

    def _inject_error(self, cursor, exc):
        self._injected_error = (cursor, exc)

    @property
    def schema(self):
        return self._schema

    def _read_next_raw_batch(self):
        if self._injected_error is not None and self._injected_error[0] <= self._pos:
            self._injected_error = None
            raise self._injected_error[1]

        if self._arrow_stream is None:
            self._arrow_stream = pa.ipc.open_stream(self._reader)

        if self._read_limit is not None and self._pos >= self._read_limit:
            warnings.warn(
                "Number of lines read via tunnel already reaches the limitation.",
                RuntimeWarning,
            )
            return None

        try:
            batch = self._arrow_stream.read_next_batch()
            self._pos += batch.num_rows
        except pa.ArrowTypeError as ex:
            if str(ex) != "Did not pass numpy.dtype object":
                raise
            else:
                raise pa.ArrowTypeError(
                    "Error caused by version mismatch. Try install numpy<1.20 or "
                    "upgrade your pyarrow version. Original message: " + str(ex)
                )
        except StopIteration:
            return None
        return batch

    def _convert_timezone(self, batch):
        from ...lib import tzlocal

        if not any(isinstance(tp, pa.TimestampType) for tp in batch.schema.types):
            return batch

        timezone = raw_timezone = options.local_timezone
        if timezone is True or timezone is None:
            timezone = str(tzlocal.get_localzone())

        cols = []
        for idx in range(batch.num_columns):
            col = batch.column(idx)
            name = batch.schema.names[idx]
            if not isinstance(col.type, pa.TimestampType):
                cols.append(col)
                continue

            if timezone is False or self._schema[name].type == types.timestamp_ntz:
                col = col.cast(pa.timestamp(col.type.unit))
            else:
                col = col.cast(pa.timestamp(col.type.unit, timezone))
                if raw_timezone is None or raw_timezone is True:
                    if hasattr(pac, "local_timestamp"):
                        col = pac.local_timestamp(col)
                    else:
                        col = pa.Array.from_pandas(
                            col.to_pandas().dt.tz_localize(None)
                        ).cast(pa.timestamp(col.type.unit))
            cols.append(col)

        return pa.RecordBatch.from_arrays(cols, names=batch.schema.names)

    def _append_partition_cols(self, batch):
        col_set = set(self._columns or [c.name for c in self._schema.columns])
        pt_obj = (
            types.PartitionSpec(self._partition_spec) if self._partition_spec else None
        )

        sel_col_set = set(self._columns or [])
        if pt_obj and any(c in sel_col_set for c in pt_obj.keys()):
            # append partitions selected in columns argument
            self._append_partitions = True
        if not pt_obj or not self._append_partitions:
            return batch

        batch_cols = list(batch.columns)
        batch_col_names = list(batch.schema.names)
        for key, val in pt_obj.items():
            if key not in col_set:
                continue
            val = types.validate_value(val, self._schema[key].type)
            batch_cols.append(pa.array(np.repeat([val], batch.num_rows)))
            batch_col_names.append(key)
        return pa.RecordBatch.from_arrays(batch_cols, names=batch_col_names)

    def read_next_batch(self):
        """
        Read next Arrow RecordBatch from tunnel.

        :return: Arrow RecordBatch
        """
        if self._reader is None:
            return None

        batch = utils.call_with_retry(
            self._read_next_raw_batch, reset_func=self._reopen_reader
        )
        if batch is None:
            return None

        batch = self._append_partition_cols(batch)
        if self._columns and self._columns != batch.schema.names:
            col_to_array = dict()
            col_name_set = set(self._columns)
            for name, arr in zip(batch.schema.names, batch.columns):
                if name not in col_name_set:
                    continue
                col_to_array[name] = arr
            arrays = [col_to_array[name] for name in self._columns]
            batch = pa.RecordBatch.from_arrays(arrays, names=self._columns)
        batch = self._convert_timezone(batch)
        return batch

    def read(self):
        """
        Read all data from tunnel and forms an Arrow Table.

        :return: Arrow Table
        """
        batches = []
        while True:
            batch = self.read_next_batch()
            if batch is None:
                break
            batches.append(batch)

        if not batches:
            return self._arrow_schema.empty_table()
        return pa.Table.from_batches(batches)

    def __iter__(self):
        return self

    def __next__(self):
        """
        Read next Arrow RecordBatch from tunnel.
        """
        batch = self.read_next_batch()
        if batch is None:
            raise StopIteration
        return batch

    @property
    def count(self):
        return self._pos

    @property
    def n_bytes(self):
        return self._last_n_bytes + len(self._reader)

    def get_total_bytes(self):
        return self.n_bytes

    def close(self):
        if hasattr(self._reader, "close"):
            self._reader.close()

    def _convert_batch_to_pandas(self, batch):
        series_list = []
        type_mapper = pd.ArrowDtype if options.tunnel.pd_cast_mode == "arrow" else None
        if not self._pd_column_converters and not self._coerce_numpy_columns:
            return batch.to_pandas(types_mapper=type_mapper)
        for col_name, arrow_column in zip(batch.schema.names, batch.columns):
            if col_name not in self._pd_column_converters:
                series = arrow_column.to_pandas(types_mapper=type_mapper)
                if col_name in self._coerce_numpy_columns and series.dtype == np.dtype(
                    "O"
                ):
                    if pa.types.is_floating(arrow_column.type):
                        col_dtype = arrow_column.type.to_pandas_dtype()
                    else:
                        col_dtype = np.dtype(float)
                    series = series.astype(col_dtype)
                series_list.append(series)
            else:
                try:
                    series = arrow_column.to_pandas(types_mapper=type_mapper)
                except pa.ArrowNotImplementedError:
                    dtype = type_mapper(arrow_column.type) if type_mapper else None
                    series = pd.Series(
                        arrow_column.to_pylist(), name=col_name, dtype=dtype
                    )
                series_list.append(series.map(self._pd_column_converters[col_name]))
        return pd.concat(series_list, axis=1)

    def to_pandas(self):
        """
        Read all data from tunnel and convert to a Pandas DataFrame.
        """
        import pandas as pd

        batches = []
        while True:
            batch = self.read_next_batch()
            if batch is None:
                break
            batches.append(self._convert_batch_to_pandas(batch))

        if not batches:
            return self._arrow_schema.empty_table().to_pandas()
        return pd.concat(batches, axis=0, ignore_index=True)

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


_reflective = lambda x: x


if convert_legacy_decimal_bytes is None:

    def convert_legacy_decimal_bytes(value):
        """
        Legacy decimal memory layout:
            int8_t  mNull;
            int8_t  mSign;
            int8_t  mIntg;
            int8_t  mFrac; only 0, 1, 2
            int32_t mData[6];
            int8_t mPadding[4]; //For Memory Align
        """
        if value is None:
            return None

        is_null, sign, intg, frac = struct.unpack("<4b", value[:4])
        if is_null:  # pragma: no cover
            return None
        if intg + frac == 0:
            return Decimal("0")

        sio = BytesIO() if compat.PY27 else StringIO()
        if sign > 0:
            sio.write("-")
        intg_nums = struct.unpack("<%dI" % intg, value[12 : 12 + intg * 4])
        intg_val = "".join("%09d" % d for d in reversed(intg_nums)).lstrip("0")
        sio.write(intg_val or "0")
        if frac > 0:
            sio.write(".")
            frac_nums = struct.unpack("<%dI" % frac, value[12 - frac * 4 : 12])
            sio.write("".join("%09d" % d for d in reversed(frac_nums)))
        return Decimal(sio.getvalue())


class ArrowRecordFieldConverter(object):
    _sensitive_types = (
        types.Datetime,
        types.Timestamp,
        types.TimestampNTZ,
        types.Array,
        types.Map,
        types.Struct,
    )

    def __init__(self, odps_type, arrow_type=None, convert_ts=True):
        self._mills_converter = utils.MillisecondsConverter()
        self._mills_converter_utc = utils.MillisecondsConverter(local_tz=False)
        self._convert_ts = convert_ts

        self._converter = self._build_converter(odps_type, arrow_type)

    def _convert_datetime(self, value):
        if value is None:
            return None
        mills = self._mills_converter.to_milliseconds(value)
        return self._mills_converter.from_milliseconds(mills)

    def _convert_ts_timestamp(self, value, ntz=False):
        if value is None:
            return None

        if not ntz:
            converter = self._mills_converter
        else:  # TimestampNtz
            converter = self._mills_converter_utc

        microsec = value.microsecond
        nanosec = value.nanosecond
        secs = int(converter.to_milliseconds(value.to_pydatetime()) / 1000)
        value = pd.Timestamp(converter.from_milliseconds(secs * 1000))
        return value.replace(microsecond=microsec, nanosecond=nanosec)

    def _convert_struct_timestamp(self, value, ntz=False):
        if value is None:
            return None

        if not ntz:
            converter = self._mills_converter
        else:  # TimestampNtz
            converter = self._mills_converter_utc

        ts = pd.Timestamp(converter.from_milliseconds(value["sec"] * 1000))
        nanos = value["nano"]
        return ts.replace(microsecond=nanos // 1000, nanosecond=nanos % 1000)

    @staticmethod
    def _convert_struct_timedelta(value):
        if value is None:
            return None

        nanos = value["nano"]
        return pd.Timedelta(
            seconds=value["sec"], microseconds=nanos // 1000, nanoseconds=nanos % 1000
        )

    @staticmethod
    def _convert_struct(value, field_converters, tuple_type, use_ordered_dict=False):
        if value is None:
            return None

        result_iter = ((k, field_converters[k](v)) for k, v in value.items())
        if tuple_type is not None:
            return tuple_type(**dict(result_iter))
        elif not use_ordered_dict:
            return dict(result_iter)
        else:
            d = dict(result_iter)
            return OrderedDict([(k, d[k]) for k in field_converters if k in d])

    def _build_converter(self, odps_type, arrow_type=None):
        import pyarrow as pa

        arrow_decimal_types = (pa.Decimal128Type,)
        if hasattr(pa, "Decimal256Type"):
            arrow_decimal_types += (pa.Decimal256Type,)

        if self._convert_ts and isinstance(odps_type, types.Datetime):
            return self._convert_datetime
        elif isinstance(odps_type, types.Timestamp):
            if isinstance(arrow_type, pa.StructType):
                return self._convert_struct_timestamp
            elif self._convert_ts:
                return self._convert_ts_timestamp
            else:
                return _reflective
        elif isinstance(odps_type, types.TimestampNTZ):
            if isinstance(arrow_type, pa.StructType):
                return functools.partial(self._convert_struct_timestamp, ntz=True)
            elif self._convert_ts:
                return functools.partial(self._convert_ts_timestamp, ntz=True)
            else:
                return _reflective
        elif (
            isinstance(odps_type, types.Decimal)
            and isinstance(arrow_type, pa.FixedSizeBinaryType)
            and not isinstance(arrow_type, arrow_decimal_types)
        ):
            return convert_legacy_decimal_bytes
        elif isinstance(odps_type, types.IntervalDayTime) and isinstance(
            arrow_type, pa.StructType
        ):
            return self._convert_struct_timedelta
        elif isinstance(odps_type, types.Array):
            arrow_value_type = getattr(arrow_type, "value_type", None)
            sub_converter = self._build_converter(
                odps_type.value_type, arrow_value_type
            )
            if sub_converter is _reflective:
                return _reflective
            return (
                lambda value: [sub_converter(x) for x in value]
                if value is not None
                else None
            )
        elif isinstance(odps_type, types.Map):
            arrow_key_type = getattr(arrow_type, "key_type", None)
            arrow_value_type = getattr(arrow_type, "item_type", None)

            key_converter = self._build_converter(odps_type.key_type, arrow_key_type)
            value_converter = self._build_converter(
                odps_type.value_type, arrow_value_type
            )
            dict_hook = OrderedDict if odps_type._use_ordered_dict else dict
            if key_converter is _reflective and value_converter is _reflective:
                return dict_hook

            return (
                lambda value: dict_hook(
                    [(key_converter(k), value_converter(v)) for k, v in value]
                )
                if value is not None
                else None
            )
        elif isinstance(odps_type, types.Struct):
            field_converters = OrderedDict()
            for field_name, field_type in odps_type.field_types.items():
                arrow_field_type = None
                if arrow_type is not None:
                    arrow_field_type = arrow_type[field_name].type
                field_converters[field_name] = self._build_converter(
                    field_type, arrow_field_type
                )

            if options.struct_as_dict:
                tuple_type = None
            else:
                tuple_type = odps_type.namedtuple_type
            use_ordered_dict = odps_type._use_ordered_dict
            return functools.partial(
                self._convert_struct,
                field_converters=field_converters,
                tuple_type=tuple_type,
                use_ordered_dict=use_ordered_dict,
            )
        else:
            return _reflective

    def __call__(self, value):
        if value is None:
            return None
        return self._converter(value)


class ArrowRecordReader(AbstractRecordReader):
    _complex_types_to_convert = (types.Array, types.Map, types.Struct)

    def __init__(self, arrow_reader, make_compat=True, read_all=False):
        self._arrow_reader = arrow_reader
        self._batch_pos = 0
        self._total_pos = 0
        self._cur_batch = None
        self._make_compat = make_compat

        self._field_converters = None

        if read_all:
            self._cur_batch = arrow_reader.read()

    def _convert_record(self, arrow_values):
        py_values = [x.as_py() for x in arrow_values]
        if not self._make_compat:
            return py_values
        else:
            return [
                converter(value)
                for value, converter in zip(py_values, self._field_converters)
            ]

    def read(self):
        if self._cur_batch is None or self._batch_pos >= self._cur_batch.num_rows:
            self._cur_batch = self._arrow_reader.read_next_batch()
            self._batch_pos = 0

            if self._cur_batch is None or self._cur_batch.num_rows == 0:
                return None

        if self._field_converters is None:
            table_schema = self._arrow_reader.schema
            self._field_converters = [
                ArrowRecordFieldConverter(table_schema[col_name].type, arrow_type)
                for col_name, arrow_type in zip(
                    self._cur_batch.schema.names, self._cur_batch.schema.types
                )
            ]
        tp = tuple(col[self._batch_pos] for col in self._cur_batch.columns)
        self._batch_pos += 1
        self._total_pos += 1

        return Record(schema=self.schema, values=self._convert_record(tp))

    def to_pandas(self, start=None, count=None, **kw):
        step = kw.get("step") or 1
        return self._arrow_reader.to_pandas().iloc[start : start + count : step]

    def close(self):
        self._arrow_reader.close()

    def __next__(self):
        rec = self.read()
        if rec is None:
            raise StopIteration
        return rec

    next = __next__

    @property
    def count(self):
        return self._total_pos

    @property
    def schema(self):
        return self._arrow_reader.schema

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()