#!/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 numpy as np import pandas as pd from mars.utils import parse_readable_size, ceildiv from mars.serialize import ( StringField, Int64Field, SeriesField, DictField, BoolField, ListField, ValueType, ) from mars.dataframe import ArrowStringDtype from mars.dataframe.operands import DataFrameOperandMixin, DataFrameOperand from mars.dataframe.utils import ( parse_index, standardize_range_index, arrow_table_to_pandas_dataframe, ) from mars.optimizes.runtime.dataframe import DataSourceHeadRule try: from mars.dataframe.datasource.core import ( HeadOptimizedDataSource, ColumnPruneSupportedDataSourceMixin, ) BASES = (HeadOptimizedDataSource, ColumnPruneSupportedDataSourceMixin) head_can_be_opt = True except ImportError: BASES = (DataFrameOperand, DataFrameOperandMixin) head_can_be_opt = False try: from mars.core import OutputType _output_type_kw = dict(_output_types=[OutputType.dataframe]) except ImportError: from mars.dataframe.operands import ObjectType _output_type_kw = dict(_object_type=ObjectType.dataframe) from ....df.backends.odpssql.types import odps_type_to_df_type from ....errors import ODPSError from ....utils import to_str, to_timestamp from ...utils import filter_partitions, check_partition_exist logger = logging.getLogger(__name__) READ_CHUNK_LIMIT = 64 * 1024**2 MAX_CHUNK_SIZE = 512 * 1024**2 ORC_COMPRESSION_RATIO = 5 STRING_FIELD_OVERHEAD = 50 _Base = type("_DataSource", BASES, dict()) class DataFrameReadTable(_Base): _op_type_ = 123450 _odps_params = DictField("odps_params") _table_name = StringField("table_name") _partition_spec = StringField("partition_spec") _dtypes = SeriesField("dtypes") _add_offset = BoolField("add_offset") _columns = ListField("columns") _nrows = Int64Field("nrows") _use_arrow_dtype = BoolField("use_arrow_dtype") _string_as_binary = BoolField("string_as_binary") _append_partitions = BoolField("append_partitions") _last_modified_time = Int64Field("last_modified_time") _with_split_meta_on_tile = BoolField("with_split_meta_on_tile") _index_columns = ListField("index_columns", ValueType.string) _index_dtypes = SeriesField("index_dtypes") def __init__( self, odps_params=None, table_name=None, partition_spec=None, columns=None, dtypes=None, nrows=None, sparse=None, add_offset=True, use_arrow_dtype=None, string_as_binary=None, memory_scale=None, append_partitions=None, last_modified_time=None, with_split_meta_on_tile=False, index_columns=None, index_dtypes=None, **kw ): kw.update(_output_type_kw) super(DataFrameReadTable, self).__init__( _odps_params=odps_params, _table_name=table_name, _partition_spec=partition_spec, _columns=columns, _dtypes=dtypes, _nrows=nrows, _sparse=sparse, _use_arrow_dtype=use_arrow_dtype, _string_as_binary=string_as_binary, _add_offset=add_offset, _append_partitions=append_partitions, _last_modified_time=last_modified_time, _memory_scale=memory_scale, _with_split_meta_on_tile=with_split_meta_on_tile, _index_columns=index_columns, _index_dtypes=index_dtypes, **kw ) @property def retryable(self): return False @property def odps_params(self): return self._odps_params @property def table_name(self): return self._table_name @property def partition(self): return getattr(self, "_partition_spec", None) @property def columns(self): return self._columns @property def dtypes(self): return self._dtypes @property def nrows(self): return self._nrows @property def use_arrow_dtype(self): return self._use_arrow_dtype @property def string_as_binary(self): return self._string_as_binary @property def add_offset(self): return self._add_offset @property def append_partitions(self): return self._append_partitions @property def with_split_meta_on_tile(self): return self._with_split_meta_on_tile @property def index_columns(self): return self._index_columns @property def index_dtypes(self): return self._index_dtypes def get_columns(self): return self._columns def set_pruned_columns(self, columns): 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): from odps import ODPS from odps.accounts import BearerTokenAccount from cupid import CupidSession, context from cupid.errors import CupidError from mars.context import get_context cupid_ctx = context() bearer_token = cupid_ctx.get_bearer_token() account = BearerTokenAccount(bearer_token) 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( None, None, account=account, project=odps_params["project"], endpoint=endpoint, ) cupid_session = CupidSession(o) mars_context = get_context() df = op.outputs[0] split_size = df.extra_params.chunk_bytes or READ_CHUNK_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) 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 = [] chunk_idx = 0 for data_src in data_srcs: try: data_store_size = data_src.size except ODPSError: # fail to get data size, just ignore pass else: if data_store_size < split_size and mars_context is not None: # get worker counts worker_count = max(len(mars_context.get_worker_addresses()), 1) # data is too small, split as many as number of cores split_size = data_store_size // worker_count # at least 1M split_size = max(split_size, 1 * 1024**2) logger.debug( "Input data size is too small, split_size is %s", split_size ) logger.debug( "Start creating download session of table %s from cupid, " "columns: %s", op.table_name, op.columns, ) while True: try: download_session = cupid_session.create_download_session( data_src, split_size=split_size, columns=op.columns, with_split_meta=op.with_split_meta_on_tile, ) break except CupidError: logger.debug( "The number of splits exceeds 100000, split_size is %s", split_size, ) if split_size >= MAX_CHUNK_SIZE: raise else: split_size *= 2 logger.debug( "%s table splits have been created.", str(len(download_session.splits)) ) meta_chunk_rows = [ split.meta_row_count for split in download_session.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 download_session.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) ] partition_spec = ( str(data_src.partition_spec) if getattr(data_src, "partition_spec", None) else None ) logger.warning("Estimated chunk rows: %r", est_chunk_rows) if len(download_session.splits) == 0: logger.debug("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(download_session.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, add_offset=op.add_offset, 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, ) # 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) if op.add_offset: 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) for data_src in data_srcs: data_store_size = data_src.size shape = out_shape 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 READ_CHUNK_LIMIT chunk_count = data_store_size // chunk_bytes + ( data_store_size % chunk_bytes != 0 ) chunk_size = ceildiv(shape[0], chunk_count) split_size = chunk_bytes else: chunk_count = ceildiv(shape[0], 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), shape[0]) 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, add_offset=op.add_offset, 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, index_columns=op.index_columns, ) 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.add_offset: 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): from cupid.runtime import RuntimeContext if RuntimeContext.is_context_ready(): return cls._tile_cupid(op) else: return cls._tile_tunnel(op) if not head_can_be_opt: tile = _tile class DataFrameReadTableSplit(_Base): _op_type_ = 123451 # for cupid _cupid_handle = StringField("cupid_handle") _split_index = Int64Field("split_index") _split_file_start = Int64Field("split_file_start") _split_file_end = Int64Field("split_file_end") _schema_file_start = Int64Field("schema_file_start") _schema_file_end = Int64Field("schema_file_end") _use_arrow_dtype = BoolField("use_arrow_dtype") _string_as_binary = BoolField("string_as_binary") _dtypes = SeriesField("dtypes") _nrows = Int64Field("nrows") # for tunnel _table_name = StringField("table_name") _partition_spec = StringField("partition_spec") _start_index = Int64Field("start_index") _end_index = Int64Field("end_index") _odps_params = DictField("odps_params") _columns = ListField("columns") _split_size = Int64Field("split_size") _append_partitions = BoolField("append_partitions") _estimate_rows = Int64Field("estimate_rows") _meta_raw_size = Int64Field("meta_raw_size") _index_columns = ListField("index_columns", ValueType.string) def __init__( self, cupid_handle=None, split_index=None, split_file_start=None, split_file_end=None, schema_file_start=None, schema_file_end=None, table_name=None, partition_spec=None, start_index=None, end_index=None, odps_params=None, columns=None, nrows=None, dtypes=None, string_as_binary=None, split_size=None, use_arrow_dtype=None, memory_scale=None, estimate_rows=None, meta_raw_size=None, append_partitions=None, sparse=None, index_columns=None, **kw ): kw.update(_output_type_kw) super(DataFrameReadTableSplit, self).__init__( _cupid_handle=cupid_handle, _split_index=split_index, _split_file_start=split_file_start, _split_file_end=split_file_end, _schema_file_start=schema_file_start, _schema_file_end=schema_file_end, _table_name=table_name, _partition_spec=partition_spec, _columns=columns, _start_index=start_index, _end_index=end_index, _odps_params=odps_params, _use_arrow_dtype=use_arrow_dtype, _string_as_binary=string_as_binary, _nrows=nrows, _estimate_rows=estimate_rows, _split_size=split_size, _dtypes=dtypes, _append_partitions=append_partitions, _sparse=sparse, _meta_raw_size=meta_raw_size, _memory_scale=memory_scale, _index_columns=index_columns, **kw ) @property def retryable(self): return False @property def output_limit(self): return 1 @property def cupid_handle(self): return self._cupid_handle @property def split_index(self): return self._split_index @property def split_file_start(self): return self._split_file_start @property def split_file_end(self): return self._split_file_end @property def schema_file_start(self): return self._schema_file_start @property def schema_file_end(self): return self._schema_file_end @property def table_name(self): return self._table_name @property def partition_spec(self): return self._partition_spec @property def start_index(self): return self._start_index @property def end_index(self): return self._end_index @property def odps_params(self): return self._odps_params @property def columns(self): return self._columns @property def nrows(self): return self._nrows @property def dtypes(self): return self._dtypes @property def split_size(self): return self._split_size @property def estimate_rows(self): return self._estimate_rows @property def use_arrow_dtype(self): return self._use_arrow_dtype @property def string_as_binary(self): return self._string_as_binary @property def append_partitions(self): return self._append_partitions @property def meta_raw_size(self): return self._meta_raw_size @property def index_columns(self): return self._index_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.debug("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.debug( "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 _execute_in_cupid(cls, ctx, op): import pyarrow as pa from cupid.io.table import TableSplit out = op.outputs[0] if op.cupid_handle is None: empty_df = pd.DataFrame() for name, dtype in out.dtypes.items(): empty_df[name] = pd.Series(dtype=dtype) ctx[out.key] = empty_df return tsp = TableSplit( _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, ) logger.debug("Read split table, split index: %s", op.split_index) reader = tsp.open_arrow_reader() if op.nrows is None: arrow_table = reader.read_all() else: nrows = 0 batches = [] while nrows < op.nrows: try: batch = reader.read_next_batch() nrows += batch.num_rows batches.append(batch) except StopIteration: break logger.debug("Read %s rows of this split.", op.nrows) arrow_table = pa.Table.from_batches(batches) arrow_table = cls._append_partition_values(arrow_table, op) if op.string_as_binary: arrow_table = cls._cast_string_to_binary(arrow_table) data = arrow_table_to_pandas_dataframe( arrow_table, use_arrow_dtype=op.use_arrow_dtype ) if op.index_columns: data = data.set_index(op.index_columns) if op.nrows is not None: data = data[: op.nrows] data = cls._align_columns(data, out.dtypes) logger.debug("Read split table finished, split index: %s", op.split_index) logger.debug( "Split data shape is %s, size is %s", data.shape, data.memory_usage(deep=True).sum(), ) ctx[out.key] = data @classmethod def _execute_arrow_tunnel(cls, ctx, op): from odps import ODPS from odps.tunnel import TableTunnel 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() is not None else None tunnel = TableTunnel(o, project=t.project) if op.partition_spec is not None: download_session = tunnel.create_download_session( t.name, partition_spec=op.partition_spec ) else: download_session = tunnel.create_download_session(t.name, schema=schema_name) logger.debug( "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 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() 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_columns(data, op.outputs[0].dtypes) logger.debug( "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): from cupid.runtime import RuntimeContext if RuntimeContext.is_context_ready(): 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=False, use_arrow_dtype=False, string_as_binary=None, memory_scale=None, append_partitions=False, with_split_meta_on_tile=False, index_columns=None, ): import pandas as pd 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) op = DataFrameReadTable( odps_params=odps_params, table_name=table.full_table_name, partition_spec=partition, dtypes=dtypes, sparse=sparse, add_offset=add_offset, 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, index_columns=index_columns, index_dtypes=index_dtypes, ) return op(shape, chunk_bytes=chunk_bytes, chunk_size=chunk_size) class ReadODPSTableHeadRule(DataSourceHeadRule): @staticmethod def match(chunk, graph, keys): from mars.dataframe.indexing.iloc import DataFrameIlocGetItem, SeriesIlocGetItem op = chunk.op inputs = graph.predecessors(chunk) if ( len(inputs) == 1 and isinstance(op, (DataFrameIlocGetItem, SeriesIlocGetItem)) and isinstance(inputs[0].op, DataFrameReadTableSplit) and inputs[0].key not in keys ): try: is_head = op.can_be_optimized() except AttributeError: is_head = op.is_head() if is_head: return True else: return False return False