odps/mars_extension/oscar/dataframe/datasource.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2022 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 logging
import os
import warnings
import time

import numpy as np
import pandas as pd

from mars.utils import parse_readable_size, ceildiv
from mars.serialization.serializables import (
    StringField,
    Int64Field,
    ListField,
    SeriesField,
    DictField,
    BoolField,
    AnyField,
    FieldTypes,
)
from mars.dataframe import ArrowStringDtype
from mars.dataframe.core import DATAFRAME_CHUNK_TYPE
from mars.dataframe.utils import (
    parse_index,
    standardize_range_index,
    arrow_table_to_pandas_dataframe,
)
from mars.dataframe.datasource.core import ColumnPruneSupportedDataSourceMixin
from mars.core import OutputType

from ....config import options
from ....df.backends.odpssql.types import odps_type_to_df_type
from ....utils import to_str, to_timestamp
from ...utils import check_partition_exist, filter_partitions
from ..cupid_service import CupidServiceClient

logger = logging.getLogger(__name__)

CHUNK_BYTES_LIMIT = 64 * 1024**2
MAX_CHUNK_NUM = 512 * 1024**2

ORC_COMPRESSION_RATIO = 5
STRING_FIELD_OVERHEAD = 50

try:
    from mars.dataframe.datasource.core import (
        IncrementalIndexDatasource,
        IncrementalIndexDataSourceMixin,
    )

    _BASE = (
        IncrementalIndexDatasource,
        ColumnPruneSupportedDataSourceMixin,
        IncrementalIndexDataSourceMixin,
    )
    _NEED_STANDARDIZE = False
except ImportError:
    from mars.dataframe.datasource.core import HeadOptimizedDataSource

    _BASE = (HeadOptimizedDataSource, ColumnPruneSupportedDataSourceMixin)
    _NEED_STANDARDIZE = True


class DataFrameReadTable(*_BASE):
    _op_type_ = 123450

    odps_params = DictField("odps_params")
    table_name = StringField("table_name")
    partition_spec = StringField("partition_spec", default=None)
    dtypes = SeriesField("dtypes", default=None)
    index_type = StringField("index_type", default=None)
    columns = AnyField("columns", default=None)
    nrows = Int64Field("nrows", default=None)
    use_arrow_dtype = BoolField("use_arrow_dtype", default=None)
    string_as_binary = BoolField("string_as_binary", default=None)
    append_partitions = BoolField("append_partitions", default=None)
    last_modified_time = Int64Field("last_modified_time", default=None)
    with_split_meta_on_tile = BoolField("with_split_meta_on_tile", default=None)
    retry_times = Int64Field("retry_times", default=None)
    tunnel_quota_name = StringField("tunnel_quota_name", default=None)
    index_columns = ListField("index_columns", FieldTypes.string, default=None)
    index_dtypes = SeriesField("index_dtypes", default=None)

    def __init__(self, sparse=None, memory_scale=None, **kw):
        super(DataFrameReadTable, self).__init__(
            sparse=sparse,
            memory_scale=memory_scale,
            _output_types=[OutputType.dataframe],
            **kw
        )

    @property
    def retryable(self):
        if "CUPID_SERVICE_SOCKET" in os.environ:
            return False
        else:
            return True

    @property
    def partition(self):
        return getattr(self, "partition_spec", None)

    @property
    def incremental_index(self):
        return self.index_type == "incremental"

    def get_columns(self):
        return self.columns

    def set_pruned_columns(self, columns, *, keep_order=None):  # pragma: no cover
        self.columns = columns

    def __call__(self, shape, chunk_bytes=None, chunk_size=None):
        import numpy as np
        import pandas as pd

        if not self.index_columns:
            if np.isnan(shape[0]):
                index_value = parse_index(pd.RangeIndex(0))
            else:
                index_value = parse_index(pd.RangeIndex(shape[0]))
        elif len(self.index_columns) == 1:
            index_value = parse_index(pd.Index([]).astype(self.index_dtypes[0]))
        else:
            idx = pd.MultiIndex.from_frame(
                pd.DataFrame([], columns=self.index_columns).astype(self.index_dtypes)
            )
            index_value = parse_index(idx)

        columns_value = parse_index(self.dtypes.index, store_data=True)
        return self.new_dataframe(
            None,
            shape,
            dtypes=self.dtypes,
            index_value=index_value,
            columns_value=columns_value,
            chunk_bytes=chunk_bytes,
            chunk_size=chunk_size,
        )

    @classmethod
    def _tile_cupid(cls, op):
        import numpy as np
        import pandas as pd
        from mars.core.context import get_context

        df = op.outputs[0]
        split_size = df.extra_params.chunk_bytes or CHUNK_BYTES_LIMIT

        out_dtypes = df.dtypes
        out_shape = df.shape
        out_columns_value = df.columns_value
        if op.columns is not None:
            out_dtypes = out_dtypes[op.columns]
            out_shape = (df.shape[0], len(op.columns))
            out_columns_value = parse_index(out_dtypes.index, store_data=True)

        mars_context = get_context()
        if mars_context is not None:
            worker_count = len(mars_context.get_worker_addresses())
        else:
            worker_count = None

        cupid_client = CupidServiceClient()
        try:
            parts = cupid_client.enum_table_partitions(
                op.odps_params, op.table_name, op.partition
            )
            if parts is None:
                parts = [None]

            out_chunks = []
            chunk_idx = 0

            for partition_spec in parts:
                columns = op.columns
                if columns is not None:
                    columns = (op.index_columns or []) + columns
                splits, split_size = cupid_client.create_table_download_session(
                    op.odps_params,
                    op.table_name,
                    partition_spec,
                    columns,
                    worker_count,
                    split_size,
                    MAX_CHUNK_NUM,
                    op.with_split_meta_on_tile,
                )

                logger.info("%s table splits have been created.", str(len(splits)))
                meta_chunk_rows = [split.meta_row_count for split in splits]
                if np.isnan(out_shape[0]):
                    est_chunk_rows = meta_chunk_rows
                else:
                    sp_file_sizes = np.array(
                        [sp.split_file_end - sp.split_file_start for sp in splits]
                    )
                    total_size = sp_file_sizes.sum()
                    ratio_chunk_rows = (
                        sp_file_sizes * out_shape[0] // total_size
                    ).tolist()
                    est_chunk_rows = [
                        mr if mr is not None else rr
                        for mr, rr in zip(meta_chunk_rows, ratio_chunk_rows)
                    ]

                logger.warning("Estimated chunk rows: %r", est_chunk_rows)

                if len(splits) == 0:
                    logger.info("Table %s has no data", op.table_name)
                    chunk_op = DataFrameReadTableSplit()
                    index_value = parse_index(pd.RangeIndex(0))
                    columns_value = parse_index(out_dtypes.index, store_data=True)
                    out_chunk = chunk_op.new_chunk(
                        None,
                        shape=(np.nan, out_shape[1]),
                        dtypes=op.dtypes,
                        index_value=index_value,
                        columns_value=columns_value,
                        index=(chunk_idx, 0),
                    )
                    out_chunks.append(out_chunk)
                    chunk_idx += 1
                else:
                    for idx, split in enumerate(splits):
                        chunk_op = DataFrameReadTableSplit(
                            cupid_handle=to_str(split.handle),
                            split_index=split.split_index,
                            split_file_start=split.split_file_start,
                            split_file_end=split.split_file_end,
                            schema_file_start=split.schema_file_start,
                            schema_file_end=split.schema_file_end,
                            index_type=op.index_type,
                            dtypes=out_dtypes,
                            sparse=op.sparse,
                            split_size=split_size,
                            string_as_binary=op.string_as_binary,
                            use_arrow_dtype=op.use_arrow_dtype,
                            estimate_rows=est_chunk_rows[idx],
                            partition_spec=partition_spec,
                            append_partitions=op.append_partitions,
                            meta_raw_size=split.meta_raw_size,
                            nrows=meta_chunk_rows[idx] or op.nrows,
                            memory_scale=op.memory_scale,
                            index_columns=op.index_columns,
                            extra_params=op.extra_params,
                        )
                        # the chunk shape is unknown
                        index_value = parse_index(pd.RangeIndex(0))
                        columns_value = parse_index(out_dtypes.index, store_data=True)
                        out_chunk = chunk_op.new_chunk(
                            None,
                            shape=(np.nan, out_shape[1]),
                            dtypes=out_dtypes,
                            index_value=index_value,
                            columns_value=columns_value,
                            index=(chunk_idx, 0),
                        )
                        chunk_idx += 1
                        out_chunks.append(out_chunk)
        finally:
            cupid_client.close()

        if op.index_type == "incremental" and _NEED_STANDARDIZE:
            out_chunks = standardize_range_index(out_chunks)
        new_op = op.copy()
        nsplits = ((np.nan,) * len(out_chunks), (out_shape[1],))
        return new_op.new_dataframes(
            None,
            shape=out_shape,
            dtypes=op.dtypes,
            index_value=df.index_value,
            columns_value=out_columns_value,
            chunks=out_chunks,
            nsplits=nsplits,
        )

    @classmethod
    def _tile_tunnel(cls, op):
        from odps import ODPS

        project = os.environ.get("ODPS_PROJECT_NAME", None)
        odps_params = op.odps_params.copy()
        if project:
            odps_params["project"] = project
        endpoint = os.environ.get("ODPS_RUNTIME_ENDPOINT") or odps_params["endpoint"]
        o = ODPS(
            odps_params["access_id"],
            odps_params["secret_access_key"],
            project=odps_params["project"],
            endpoint=endpoint,
        )

        table_obj = o.get_table(op.table_name)
        if not table_obj.table_schema.partitions:
            data_srcs = [table_obj]
        elif op.partition is not None and check_partition_exist(
            table_obj, op.partition
        ):
            data_srcs = [table_obj.get_partition(op.partition)]
        else:
            data_srcs = list(table_obj.partitions)
            if op.partition is not None:
                data_srcs = filter_partitions(o, data_srcs, op.partition)

        out_chunks = []
        row_nsplits = []
        index_start = 0
        df = op.outputs[0]

        out_dtypes = df.dtypes
        out_shape = df.shape
        out_columns_value = df.columns_value
        if op.columns is not None:
            out_dtypes = out_dtypes[op.columns]
            out_shape = (df.shape[0], len(op.columns))
            out_columns_value = parse_index(out_dtypes.index, store_data=True)

        if len(data_srcs) == 0:
            # no partitions are selected
            chunk_op = DataFrameReadTableSplit()
            index_value = parse_index(pd.RangeIndex(0))
            columns_value = parse_index(out_dtypes.index, store_data=True)
            out_chunk = chunk_op.new_chunk(
                None,
                shape=(0, out_shape[1]),
                dtypes=op.dtypes,
                index_value=index_value,
                columns_value=columns_value,
                index=(index_start, 0),
            )
            out_chunks.append(out_chunk)
        else:
            retry_times = op.retry_times or options.retry_times
            for data_src in data_srcs:
                data_store_size = data_src.size

                retries = 0
                while True:
                    try:
                        with data_src.open_reader() as reader:
                            record_count = reader.count
                        break
                    except:
                        if retries >= retry_times:
                            raise
                        retries += 1
                        time.sleep(1)
                if data_store_size == 0:
                    # empty table
                    chunk_op = DataFrameReadTableSplit()
                    index_value = parse_index(pd.RangeIndex(0))
                    columns_value = parse_index(out_dtypes.index, store_data=True)
                    out_chunk = chunk_op.new_chunk(
                        None,
                        shape=(0, out_shape[1]),
                        dtypes=op.dtypes,
                        index_value=index_value,
                        columns_value=columns_value,
                        index=(index_start, 0),
                    )
                    out_chunks.append(out_chunk)
                    index_start += 1
                    continue
                chunk_size = df.extra_params.chunk_size

                partition_spec = (
                    str(data_src.partition_spec)
                    if getattr(data_src, "partition_spec", None)
                    else None
                )

                if chunk_size is None:
                    chunk_bytes = df.extra_params.chunk_bytes or CHUNK_BYTES_LIMIT
                    chunk_count = data_store_size // chunk_bytes + (
                        data_store_size % chunk_bytes != 0
                    )
                    chunk_size = ceildiv(record_count, chunk_count)
                    split_size = chunk_bytes
                else:
                    chunk_count = ceildiv(record_count, chunk_size)
                    split_size = data_store_size // chunk_count

                for i in range(chunk_count):
                    start_index = chunk_size * i
                    end_index = min(chunk_size * (i + 1), record_count)
                    row_size = end_index - start_index
                    chunk_op = DataFrameReadTableSplit(
                        table_name=op.table_name,
                        partition_spec=partition_spec,
                        start_index=start_index,
                        end_index=end_index,
                        nrows=op.nrows,
                        odps_params=op.odps_params,
                        columns=op.columns,
                        index_type=op.index_type,
                        dtypes=out_dtypes,
                        sparse=op.sparse,
                        split_size=split_size,
                        use_arrow_dtype=op.use_arrow_dtype,
                        estimate_rows=row_size,
                        append_partitions=op.append_partitions,
                        memory_scale=op.memory_scale,
                        retry_times=op.retry_times,
                        tunnel_quota_name=op.tunnel_quota_name,
                        index_columns=op.index_columns,
                        extra_params=op.extra_params,
                    )
                    index_value = parse_index(pd.RangeIndex(start_index, end_index))
                    columns_value = parse_index(out_dtypes.index, store_data=True)
                    out_chunk = chunk_op.new_chunk(
                        None,
                        shape=(row_size, out_shape[1]),
                        dtypes=out_dtypes,
                        index_value=index_value,
                        columns_value=columns_value,
                        index=(index_start + i, 0),
                    )
                    row_nsplits.append(row_size)
                    out_chunks.append(out_chunk)

                index_start += chunk_count

        if op.index_type == "incremental" and _NEED_STANDARDIZE:
            out_chunks = standardize_range_index(out_chunks)

        new_op = op.copy()
        nsplits = (tuple(row_nsplits), (out_shape[1],))
        return new_op.new_dataframes(
            None,
            shape=out_shape,
            dtypes=op.dtypes,
            index_value=df.index_value,
            columns_value=out_columns_value,
            chunks=out_chunks,
            nsplits=nsplits,
        )

    @classmethod
    def _tile(cls, op):
        if "CUPID_SERVICE_SOCKET" in os.environ:
            return cls._tile_cupid(op)
        else:
            return cls._tile_tunnel(op)


class DataFrameReadTableSplit(*_BASE):
    _op_type_ = 123451

    # for cupid
    cupid_handle = StringField("cupid_handle", default=None)
    split_index = Int64Field("split_index", default=None)
    split_file_start = Int64Field("split_file_start", default=None)
    split_file_end = Int64Field("split_file_end", default=None)
    schema_file_start = Int64Field("schema_file_start", default=None)
    schema_file_end = Int64Field("schema_file_end", default=None)
    use_arrow_dtype = BoolField("use_arrow_dtype", default=None)
    string_as_binary = BoolField("string_as_binary", default=None)
    dtypes = SeriesField("dtypes", default=None)
    nrows = Int64Field("nrows", default=None)
    index_type = StringField("index_type", default=None)

    # for tunnel
    table_name = StringField("table_name", default=None)
    partition_spec = StringField("partition_spec", default=None)
    start_index = Int64Field("start_index", default=None)
    end_index = Int64Field("end_index", default=None)
    odps_params = DictField("odps_params", default=None)
    columns = AnyField("columns", default=None)
    tunnel_quota_name = StringField("tunnel_quota_name", default=None)

    split_size = Int64Field("split_size", default=None)
    append_partitions = BoolField("append_partitions", default=None)
    estimate_rows = Int64Field("estimate_rows", default=None)
    meta_raw_size = Int64Field("meta_raw_size", default=None)
    retry_times = Int64Field("retry_times", default=None)
    index_columns = ListField("index_columns", FieldTypes.string, default=None)

    def __init__(self, memory_scale=None, sparse=None, **kw):
        super(DataFrameReadTableSplit, self).__init__(
            sparse=sparse,
            memory_scale=memory_scale,
            _output_types=[OutputType.dataframe],
            **kw
        )

    @property
    def retryable(self):
        if "CUPID_SERVICE_SOCKET" in os.environ:
            return False
        else:
            return True

    @property
    def output_limit(self):
        return 1

    @property
    def incremental_index(self):
        return self.index_type == "incremental"

    def set_pruned_columns(self, columns, *, keep_order=None):  # pragma: no cover
        if isinstance(columns, str):
            columns = [columns]
        self.columns = list(columns)

    def get_columns(self):
        return self.columns

    @classmethod
    def estimate_size(cls, ctx, op):
        import numpy as np

        def is_object_dtype(dtype):
            try:
                return (
                    np.issubdtype(dtype, np.object_)
                    or np.issubdtype(dtype, np.unicode_)
                    or np.issubdtype(dtype, np.bytes_)
                )
            except TypeError:  # pragma: no cover
                return False

        if op.split_size is None:
            ctx[op.outputs[0].key] = (0, 0)
            return

        arrow_size = (op.memory_scale or ORC_COMPRESSION_RATIO) * op.split_size
        if op.meta_raw_size is not None:
            raw_arrow_size = (op.memory_scale or 1) * op.meta_raw_size
            arrow_size = max(arrow_size, raw_arrow_size)

        n_strings = len([dt for dt in op.dtypes if is_object_dtype(dt)])
        if op.estimate_rows or op.nrows:
            rows = op.nrows if op.nrows is not None else op.estimate_rows
            pd_size = arrow_size + n_strings * rows * STRING_FIELD_OVERHEAD
            logger.info("Estimate pandas memory cost: %r", pd_size)
        else:
            pd_size = arrow_size * 10 if n_strings else arrow_size

        ctx[op.outputs[0].key] = (pd_size, pd_size + arrow_size)

    @classmethod
    def _cast_string_to_binary(cls, arrow_table):
        import pyarrow as pa

        new_schema = []
        for field in arrow_table.schema:
            if field.type == pa.string():
                new_schema.append(pa.field(field.name, pa.binary()))
            else:
                new_schema.append(field)

        return arrow_table.cast(pa.schema(new_schema))

    @classmethod
    def _append_partition_values(cls, arrow_table, op):
        import pyarrow as pa

        if op.append_partitions and op.partition_spec:
            from odps.types import PartitionSpec

            spec = PartitionSpec(op.partition_spec)

            for col_name, pt_val in spec.items():
                arrow_table = arrow_table.append_column(
                    col_name, pa.array([pt_val] * arrow_table.num_rows, pa.string())
                )

        return arrow_table

    @staticmethod
    def _align_columns(data, expected_dtypes):
        data_columns = data.dtypes.index
        expected_columns = expected_dtypes.index
        if not data_columns.equals(expected_columns):
            logger.info(
                "Data columns differs from output columns, "
                "data columns: %s, output columns: %s",
                data_columns,
                expected_columns,
            )
            data.columns = expected_columns[: len(data.columns)]
            for extra_col in expected_columns[len(data.columns) :]:
                data[extra_col] = pd.Series([], dtype=expected_dtypes[extra_col])
            if not data.dtypes.index.equals(expected_columns):
                data = data[expected_columns]
        return data

    @classmethod
    def _align_output_data(cls, op, data):
        if isinstance(op.outputs[0], DATAFRAME_CHUNK_TYPE):
            dtypes = op.outputs[0].dtypes
            data = cls._align_columns(data, dtypes)
        else:
            dtypes = pd.Series([op.outputs[0].dtype], index=[op.outputs[0].name])
            data = cls._align_columns(data, dtypes)
            data = data[op.outputs[0].name]
        return data

    @classmethod
    def _build_empty_df(cls, out):
        empty_df = pd.DataFrame()
        for name, dtype in out.dtypes.items():
            empty_df[name] = pd.Series(dtype=dtype)
        return empty_df

    @classmethod
    def _execute_in_cupid(cls, ctx, op):
        out = op.outputs[0]

        if op.cupid_handle is None:
            ctx[out.key] = cls._build_empty_df(out)
            return

        split_config = dict(
            _handle=op.cupid_handle,
            _split_index=op.split_index,
            _split_file_start=op.split_file_start,
            _split_file_end=op.split_file_end,
            _schema_file_start=op.schema_file_start,
            _schema_file_end=op.schema_file_end,
        )
        cupid_client = CupidServiceClient()
        try:
            pa_table = cupid_client.read_table_data(split_config, op.nrows)
        finally:
            cupid_client.close()
            cupid_client = None
        pa_table = cls._append_partition_values(pa_table, op)

        if op.string_as_binary:
            pa_table = cls._cast_string_to_binary(pa_table)
        data = arrow_table_to_pandas_dataframe(
            pa_table, use_arrow_dtype=op.use_arrow_dtype
        )[: op.nrows]
        if op.index_columns:
            data = data.set_index(op.index_columns)

        data = cls._align_output_data(op, data)
        logger.info("Read split table finished, split index: %s", op.split_index)
        logger.info("Split data shape is %s", data.shape)
        ctx[out.key] = data

    @classmethod
    def _execute_arrow_tunnel(cls, ctx, op):
        from odps import ODPS
        from odps.tunnel import TableTunnel

        out = op.outputs[0]

        if op.table_name is None:
            # is empty table
            ctx[out.key] = cls._build_empty_df(out)
            return

        project = os.environ.get("ODPS_PROJECT_NAME", None)
        odps_params = op.odps_params.copy()
        if project:
            odps_params["project"] = project
        endpoint = os.environ.get("ODPS_RUNTIME_ENDPOINT") or odps_params["endpoint"]
        o = ODPS(
            odps_params["access_id"],
            odps_params["secret_access_key"],
            project=odps_params["project"],
            endpoint=endpoint,
        )

        t = o.get_table(op.table_name)
        schema_name = t.get_schema().name if t.get_schema() else None
        tunnel = TableTunnel(o, project=t.project, quota_name=op.tunnel_quota_name)
        retry_times = op.retry_times or options.retry_times
        init_sleep_secs = 1
        logger.info(
            "Start creating download session for table %s(%s) start index %s end index %s retry_times %s.",
            op.table_name,
            op.partition_spec,
            op.start_index,
            op.end_index,
            retry_times,
        )
        retries = 0
        while True:
            try:
                if op.partition_spec is not None:
                    download_session = tunnel.create_download_session(
                        t.name,
                        partition_spec=op.partition_spec,
                        schema=schema_name,
                        async_mode=True,
                    )
                else:
                    download_session = tunnel.create_download_session(
                        t.name, async_mode=True
                    )
                break
            except:
                if retries >= retry_times:
                    raise
                retries += 1
                sleep_secs = retries * init_sleep_secs
                logger.exception(
                    "Create download session failed, sleep %s seconds and retry it",
                    sleep_secs,
                    exc_info=1,
                )
                time.sleep(sleep_secs)

        logger.info(
            "Start reading table %s(%s) split from %s to %s",
            op.table_name,
            op.partition_spec,
            op.start_index,
            op.end_index,
        )
        if op.nrows is None:
            count = op.end_index - op.start_index
        else:
            count = op.nrows

        retries = 0
        while True:
            try:
                columns = op.columns
                if columns is not None:
                    columns = (op.index_columns or []) + op.columns
                with download_session.open_arrow_reader(
                    op.start_index, count, columns=columns
                ) as reader:
                    table = reader.read()
                break
            except:
                if retries >= retry_times:
                    raise
                retries += 1
                sleep_secs = retries * init_sleep_secs
                logger.exception(
                    "Read table failed, sleep %s seconds and retry it",
                    sleep_secs,
                    exc_info=1,
                )
                time.sleep(sleep_secs)

        table = cls._append_partition_values(table, op)
        if op.string_as_binary:
            table = cls._cast_string_to_binary(table)
        data = arrow_table_to_pandas_dataframe(
            table, use_arrow_dtype=op.use_arrow_dtype
        )
        if op.index_columns:
            data = data.set_index(op.index_columns)
        data = cls._align_output_data(op, data)
        logger.info(
            "Finish reading table %s(%s) split from %s to %s",
            op.table_name,
            op.partition_spec,
            op.start_index,
            op.end_index,
        )
        ctx[op.outputs[0].key] = data

    @classmethod
    def execute(cls, ctx, op):
        if "CUPID_SERVICE_SOCKET" in os.environ:
            cls._execute_in_cupid(ctx, op)
        else:
            cls._execute_arrow_tunnel(ctx, op)


def df_type_to_np_type(df_type, use_arrow_dtype=False):
    from ....df import types
    from ....df.backends.pd.types import _df_to_np_types

    if df_type == types.string:
        if use_arrow_dtype:
            return ArrowStringDtype()
        else:
            return np.dtype("object")
    elif df_type in _df_to_np_types:
        return _df_to_np_types[df_type]
    elif df_type == types.timestamp:
        return np.datetime64(0, "ns").dtype
    else:
        return np.dtype("object")


def read_odps_table(
    table,
    shape,
    partition=None,
    sparse=False,
    chunk_bytes=None,
    chunk_size=None,
    columns=None,
    odps_params=None,
    add_offset=None,
    use_arrow_dtype=False,
    string_as_binary=None,
    memory_scale=None,
    append_partitions=False,
    with_split_meta_on_tile=False,
    index_type="incremental",
    tunnel_quota_name=None,
    index_columns=None,
    extra_params=None,
):
    import pandas as pd

    if add_offset is not None:
        warnings.warn(
            "add_offset is deprecated, please use index_type instead",
            DeprecationWarning,
        )
        if add_offset in (True, False):
            index_type = "incremental" if add_offset else "chunk_incremental"

    if isinstance(chunk_size, (list, tuple)):
        if len(chunk_size) == 1:
            chunk_size = chunk_size[0]
        if len(chunk_size) > 1:
            raise ValueError("Only support split on rows")

    if chunk_bytes is not None:
        chunk_bytes = int(parse_readable_size(chunk_bytes)[0])
    cols = table.table_schema.columns if append_partitions else table.table_schema.simple_columns
    table_columns = [c.name for c in cols]
    table_types = [c.type for c in cols]
    df_types = [
        df_type_to_np_type(odps_type_to_df_type(type), use_arrow_dtype=use_arrow_dtype)
        for type in table_types
    ]

    if isinstance(index_columns, str):
        index_columns = [index_columns]
    if index_columns and columns is None:
        index_col_set = set(index_columns)
        columns = [c for c in table_columns if c not in index_col_set]

    if not index_columns:
        index_dtypes = None
    else:
        table_index_types = [
            df_types[table_columns.index(col)] for col in index_columns
        ]
        index_dtypes = pd.Series(table_index_types, index=index_columns)

    if columns is not None:
        table_col_set = set(columns)
        if any(col in table_col_set for col in index_columns):
            raise ValueError("Index columns and columns shall not overlap.")

        # reorder columns
        new_columns = [c for c in table_columns if c in columns]
        df_types = [df_types[table_columns.index(col)] for col in new_columns]
        table_columns = new_columns
        columns = new_columns

    dtypes = pd.Series(df_types, index=table_columns)
    retry_times = options.retry_times
    op = DataFrameReadTable(
        odps_params=odps_params,
        table_name=table.full_table_name,
        partition_spec=partition,
        dtypes=dtypes,
        sparse=sparse,
        index_type=index_type,
        columns=columns,
        use_arrow_dtype=use_arrow_dtype,
        string_as_binary=string_as_binary,
        memory_scale=memory_scale,
        append_partitions=append_partitions,
        last_modified_time=to_timestamp(table.last_data_modified_time),
        with_split_meta_on_tile=with_split_meta_on_tile,
        retry_times=retry_times,
        tunnel_quota_name=tunnel_quota_name,
        index_columns=index_columns,
        index_dtypes=index_dtypes,
        extra_params=extra_params or dict(),
    )
    return op(shape, chunk_bytes=chunk_bytes, chunk_size=chunk_size)