maxframe 0.1.0b5__cp310-cp310-macosx_10_9_universal2.whl → 1.0.0rc2__cp310-cp310-macosx_10_9_universal2.whl

maxframe/odpsio/tableio.py

@@ -16,7 +16,7 @@ import os
 import time
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import pyarrow as pa
 from odps import ODPS
@@ -25,8 +25,16 @@ from odps.apis.storage_api import (
     TableBatchScanResponse,
     TableBatchWriteResponse,
 )
-from odps.types import PartitionSpec
+from odps.config import option_context as pyodps_option_context
+from odps.tunnel import TableTunnel
+from odps.types import OdpsSchema, PartitionSpec, timestamp_ntz
 
+try:
+    import pyarrow.compute as pac
+except ImportError:
+    pac = None
+
+from ..config import options
 from ..env import ODPS_STORAGE_API_ENDPOINT
 from .schema import odps_schema_to_arrow_schema
 
@@ -35,18 +43,55 @@ PartitionsType = Union[List[str], str, None]
 _DEFAULT_ROW_BATCH_SIZE = 4096
 
 
-class MCTableIO(ABC):
+@contextmanager
+def _sync_pyodps_timezone():
+    with pyodps_option_context() as cfg:
+        cfg.local_timezone = options.local_timezone
+        yield
+
+
+class ODPSTableIO(ABC):
+    def __new__(cls, odps: ODPS):
+        if cls is ODPSTableIO:
+            if options.use_common_table:
+                return HaloTableIO(odps)
+            else:
+                return TunnelTableIO(odps)
+        return super().__new__(cls)
+
     def __init__(self, odps: ODPS):
         self._odps = odps
 
+    @classmethod
+    def _get_reader_schema(
+        cls,
+        table_schema: OdpsSchema,
+        columns: Optional[List[str]] = None,
+        partition_columns: Union[None, bool, List[str]] = None,
+    ) -> OdpsSchema:
+        final_cols = []
+
+        columns = columns or [col.name for col in table_schema.simple_columns]
+        if partition_columns is True:
+            partition_columns = [c.name for c in table_schema.partitions]
+        else:
+            partition_columns = partition_columns or []
+
+        for col_name in columns + partition_columns:
+            final_cols.append(table_schema[col_name])
+        return OdpsSchema(final_cols)
+
     @abstractmethod
     def open_reader(
         self,
         full_table_name: str,
         partitions: PartitionsType = None,
         columns: Optional[List[str]] = None,
+        partition_columns: Union[None, bool, List[str]] = None,
         start: Optional[int] = None,
-        count: Optional[int] = None,
+        stop: Optional[int] = None,
+        reverse_range: bool = False,
+        row_batch_size: int = _DEFAULT_ROW_BATCH_SIZE,
     ):
         raise NotImplementedError
 
@@ -60,12 +105,264 @@ class MCTableIO(ABC):
         raise NotImplementedError
 
 
+class TunnelMultiPartitionReader:
+    def __init__(
+        self,
+        odps_entry: ODPS,
+        table_name: str,
+        partitions: PartitionsType,
+        columns: Optional[List[str]] = None,
+        partition_columns: Optional[List[str]] = None,
+        start: Optional[int] = None,
+        count: Optional[int] = None,
+        partition_to_download_ids: Dict[str, str] = None,
+    ):
+        self._odps_entry = odps_entry
+        self._table = odps_entry.get_table(table_name)
+        self._columns = columns
+
+        odps_schema = ODPSTableIO._get_reader_schema(
+            self._table.table_schema, columns, partition_columns
+        )
+        self._schema = odps_schema_to_arrow_schema(odps_schema)
+
+        self._start = start or 0
+        self._count = count
+        self._row_left = count
+
+        self._cur_reader = None
+        self._reader_iter = None
+        self._cur_partition_id = -1
+        self._reader_start_pos = 0
+
+        if partitions is None or isinstance(partitions, str):
+            self._partitions = [partitions]
+        else:
+            self._partitions = partitions
+
+        self._partition_cols = partition_columns
+        self._partition_to_download_ids = partition_to_download_ids or dict()
+
+    @property
+    def count(self) -> Optional[int]:
+        if len(self._partitions) > 1:
+            return None
+        return self._count
+
+    def _open_next_reader(self):
+        if self._cur_reader is not None:
+            self._reader_start_pos += self._cur_reader.count
+
+        if (
+            self._row_left is not None and self._row_left <= 0
+        ) or 1 + self._cur_partition_id >= len(self._partitions):
+            self._cur_reader = None
+            return
+
+        while 1 + self._cur_partition_id < len(self._partitions):
+            self._cur_partition_id += 1
+
+            part_str = self._partitions[self._cur_partition_id]
+            with _sync_pyodps_timezone():
+                self._cur_reader = self._table.open_reader(
+                    part_str,
+                    columns=self._columns,
+                    arrow=True,
+                    download_id=self._partition_to_download_ids.get(part_str),
+                )
+            if self._cur_reader.count + self._reader_start_pos > self._start:
+                start = self._start - self._reader_start_pos
+                if self._count is None:
+                    count = None
+                else:
+                    count = min(self._count, self._cur_reader.count - start)
+
+                with _sync_pyodps_timezone():
+                    self._reader_iter = self._cur_reader.read(start, count)
+                break
+            self._reader_start_pos += self._cur_reader.count
+        else:
+            self._cur_reader = None
+
+    def _fill_batch_partition(self, batch: pa.RecordBatch) -> pa.RecordBatch:
+        pt_spec = PartitionSpec(self._partitions[self._cur_partition_id])
+
+        names = list(batch.schema.names)
+        arrays = []
+        for idx in range(batch.num_columns):
+            col = batch.column(idx)
+            if isinstance(col.type, pa.TimestampType):
+                if col.type.tz is not None:
+                    target_type = pa.timestamp(
+                        self._schema.types[idx].unit, col.type.tz
+                    )
+                    arrays.append(col.cast(target_type))
+                else:
+                    target_type = pa.timestamp(
+                        self._schema.types[idx].unit, options.local_timezone
+                    )
+                    pd_col = col.to_pandas().dt.tz_localize(options.local_timezone)
+                    arrays.append(pa.Array.from_pandas(pd_col).cast(target_type))
+            else:
+                arrays.append(batch.column(idx))
+
+        for part_col in self._partition_cols or []:
+            names.append(part_col)
+            col_type = self._schema.field_by_name(part_col).type
+            arrays.append(pa.array([pt_spec[part_col]] * batch.num_rows).cast(col_type))
+        return pa.RecordBatch.from_arrays(arrays, names)
+
+    def read(self):
+        with _sync_pyodps_timezone():
+            if self._cur_reader is None:
+                self._open_next_reader()
+                if self._cur_reader is None:
+                    return None
+            while self._cur_reader is not None:
+                try:
+                    batch = next(self._reader_iter)
+                    if batch is not None:
+                        if self._row_left is not None:
+                            self._row_left -= batch.num_rows
+                        return self._fill_batch_partition(batch)
+                except StopIteration:
+                    self._open_next_reader()
+            return None
+
+    def read_all(self) -> pa.Table:
+        batches = []
+        while True:
+            batch = self.read()
+            if batch is None:
+                break
+            batches.append(batch)
+        if not batches:
+            return self._schema.empty_table()
+        return pa.Table.from_batches(batches)
+
+
+class TunnelWrappedWriter:
+    def __init__(self, nested_writer):
+        self._writer = nested_writer
+
+    def write(self, data: Union[pa.RecordBatch, pa.Table]):
+        if not any(isinstance(tp, pa.TimestampType) for tp in data.schema.types):
+            self._writer.write(data)
+            return
+        pa_type = type(data)
+        arrays = []
+        for idx in range(data.num_columns):
+            name = data.schema.names[idx]
+            col = data.column(idx)
+            if not isinstance(col.type, pa.TimestampType):
+                arrays.append(col)
+                continue
+            if self._writer.schema[name].type == timestamp_ntz:
+                col = HaloTableArrowWriter._localize_timezone(col, "UTC")
+            else:
+                col = HaloTableArrowWriter._localize_timezone(col)
+            arrays.append(col)
+        data = pa_type.from_arrays(arrays, names=data.schema.names)
+        self._writer.write(data)
+
+    def __getattr__(self, item):
+        return getattr(self._writer, item)
+
+
+class TunnelTableIO(ODPSTableIO):
+    @contextmanager
+    def open_reader(
+        self,
+        full_table_name: str,
+        partitions: PartitionsType = None,
+        columns: Optional[List[str]] = None,
+        partition_columns: Union[None, bool, List[str]] = None,
+        start: Optional[int] = None,
+        stop: Optional[int] = None,
+        reverse_range: bool = False,
+        row_batch_size: int = _DEFAULT_ROW_BATCH_SIZE,
+    ):
+        table = self._odps.get_table(full_table_name)
+        if partition_columns is True:
+            partition_columns = [c.name for c in table.table_schema.partitions]
+
+        total_records = None
+        part_to_down_id = None
+        if (
+            (start is not None and start < 0)
+            or (stop is not None and stop < 0)
+            or (reverse_range and start is None)
+        ):
+            table = self._odps.get_table(full_table_name)
+            tunnel = TableTunnel(self._odps)
+            parts = (
+                [partitions]
+                if partitions is None or isinstance(partitions, str)
+                else partitions
+            )
+            part_to_down_id = dict()
+            total_records = 0
+            for part in parts:
+                down_session = tunnel.create_download_session(
+                    table, async_mode=True, partition_spec=part
+                )
+                part_to_down_id[part] = down_session.id
+                total_records += down_session.count
+
+        count = None
+        if start is not None or stop is not None:
+            if reverse_range:
+                start = start if start is not None else total_records - 1
+                stop = stop if stop is not None else -1
+            else:
+                start = start if start is not None else 0
+                stop = stop if stop is not None else None
+            start = start if start >= 0 else total_records + start
+            stop = stop if stop is None or stop >= 0 else total_records + stop
+            if reverse_range:
+                count = start - stop
+                start = stop + 1
+            else:
+                count = stop - start if stop is not None and start is not None else None
+
+        yield TunnelMultiPartitionReader(
+            self._odps,
+            full_table_name,
+            partitions=partitions,
+            columns=columns,
+            partition_columns=partition_columns,
+            start=start,
+            count=count,
+            partition_to_download_ids=part_to_down_id,
+        )
+
+    @contextmanager
+    def open_writer(
+        self,
+        full_table_name: str,
+        partition: Optional[str] = None,
+        overwrite: bool = True,
+    ):
+        table = self._odps.get_table(full_table_name)
+        with _sync_pyodps_timezone():
+            with table.open_writer(
+                partition=partition,
+                arrow=True,
+                create_partition=partition is not None,
+                overwrite=overwrite,
+            ) as writer:
+                # fixme should yield writer directly once pyodps fixes
+                #  related arrow timestamp bug when provided schema and
+                #  table schema is identical.
+                yield TunnelWrappedWriter(writer)
+
+
 class HaloTableArrowReader:
     def __init__(
         self,
         client: StorageApiArrowClient,
         scan_info: TableBatchScanResponse,
-        schema: pa.Schema,
+        odps_schema: OdpsSchema,
         start: Optional[int] = None,
         count: Optional[int] = None,
         row_batch_size: Optional[int] = None,
@@ -76,7 +373,8 @@ class HaloTableArrowReader:
         self._cur_split_id = -1
         self._cur_reader = None
 
-        self._schema = schema
+        self._odps_schema = odps_schema
+        self._arrow_schema = odps_schema_to_arrow_schema(odps_schema)
 
         self._start = start
         self._count = count
@@ -115,6 +413,34 @@ class HaloTableArrowReader:
         self._cur_reader = self._client.read_rows_arrow(req)
         self._cur_split_id += 1
 
+    def _convert_timezone(self, batch: pa.RecordBatch) -> pa.RecordBatch:
+        timezone = options.local_timezone
+        if not any(isinstance(tp, pa.TimestampType) for tp in batch.schema.types):
+            return batch
+
+        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 self._odps_schema[name].type == timestamp_ntz:
+                col = col.cast(pa.timestamp(col.type.unit))
+                cols.append(col)
+                continue
+
+            if hasattr(pac, "local_timestamp"):
+                col = col.cast(pa.timestamp(col.type.unit, timezone))
+            else:
+                pd_col = col.to_pandas().dt.tz_convert(timezone)
+                col = pa.Array.from_pandas(pd_col).cast(
+                    pa.timestamp(col.type.unit, timezone)
+                )
+            cols.append(col)
+
+        return pa.RecordBatch.from_arrays(cols, names=batch.schema.names)
+
     def read(self):
         if self._cur_reader is None:
             self._open_next_reader()
@@ -123,7 +449,7 @@ class HaloTableArrowReader:
         while self._cur_reader is not None:
             batch = self._cur_reader.read()
             if batch is not None:
-                return batch
+                return self._convert_timezone(batch)
             self._open_next_reader()
         return None
 
@@ -135,16 +461,21 @@ class HaloTableArrowReader:
                 break
             batches.append(batch)
         if not batches:
-            return self._schema.empty_table()
+            return self._arrow_schema.empty_table()
         return pa.Table.from_batches(batches)
 
 
 class HaloTableArrowWriter:
     def __init__(
-        self, client: StorageApiArrowClient, write_info: TableBatchWriteResponse
+        self,
+        client: StorageApiArrowClient,
+        write_info: TableBatchWriteResponse,
+        odps_schema: OdpsSchema,
     ):
         self._client = client
         self._write_info = write_info
+        self._odps_schema = odps_schema
+        self._arrow_schema = odps_schema_to_arrow_schema(odps_schema)
 
         self._writer = None
 
@@ -155,12 +486,52 @@ class HaloTableArrowWriter:
             WriteRowsRequest(self._write_info.session_id)
         )
 
+    @classmethod
+    def _localize_timezone(cls, col, tz=None):
+        from odps.lib import tzlocal
+
+        if tz is None:
+            if options.local_timezone is None:
+                tz = str(tzlocal.get_localzone())
+            else:
+                tz = str(options.local_timezone)
+
+        if col.type.tz is not None:
+            return col
+        if hasattr(pac, "assume_timezone"):
+            col = pac.assume_timezone(col, tz)
+            return col
+        else:
+            col = col.to_pandas()
+            return pa.Array.from_pandas(col.dt.tz_localize(tz))
+
+    def _convert_schema(self, batch: pa.RecordBatch):
+        if batch.schema == self._arrow_schema and not any(
+            isinstance(tp, pa.TimestampType) for tp in self._arrow_schema.types
+        ):
+            return batch
+        cols = []
+        for idx in range(batch.num_columns):
+            col = batch.column(idx)
+            name = batch.schema.names[idx]
+
+            if isinstance(col.type, pa.TimestampType):
+                if self._odps_schema[name].type == timestamp_ntz:
+                    col = self._localize_timezone(col, "UTC")
+                else:
+                    col = self._localize_timezone(col)
+
+            if col.type != self._arrow_schema.types[idx]:
+                col = col.cast(self._arrow_schema.types[idx])
+            cols.append(col)
+        return pa.RecordBatch.from_arrays(cols, names=batch.schema.names)
+
     def write(self, batch):
         if isinstance(batch, pa.Table):
             for b in batch.to_batches():
-                self._writer.write(b)
+                self._writer.write(self._convert_schema(b))
         else:
-            self._writer.write(batch)
+            self._writer.write(self._convert_schema(batch))
 
     def close(self):
         commit_msg, is_success = self._writer.finish()
@@ -169,7 +540,7 @@ class HaloTableArrowWriter:
         return commit_msg
 
 
-class HaloTableIO(MCTableIO):
+class HaloTableIO(ODPSTableIO):
     _storage_api_endpoint = os.getenv(ODPS_STORAGE_API_ENDPOINT)
 
     @staticmethod
@@ -275,10 +646,13 @@ class HaloTableIO(MCTableIO):
             else:
                 count = stop - start
 
+        reader_schema = self._get_reader_schema(
+            table.table_schema, columns, partition_columns
+        )
         yield HaloTableArrowReader(
             client,
             resp,
-            schema=odps_schema_to_arrow_schema(table.table_schema),
+            odps_schema=reader_schema,
             start=start,
             count=count,
             row_batch_size=row_batch_size,
@@ -308,7 +682,7 @@ class HaloTableIO(MCTableIO):
         resp = client.create_write_session(req)
 
         session_id = resp.session_id
-        writer = HaloTableArrowWriter(client, resp)
+        writer = HaloTableArrowWriter(client, resp, table.table_schema)
         writer.open()
 
         yield writer

maxframe/odpsio/tests/test_schema.py

@@ -143,17 +143,17 @@ def test_pandas_to_odps_schema_index(wrap_obj):
     data = pd.Index(np.random.randint(0, 100, 100))
 
     test_idx = _wrap_maxframe_obj(data, wrap=wrap_obj)
-    schema, meta = pandas_to_odps_schema(test_idx, unknown_as_string=True)
-    assert [c.name for c in schema.columns] == ["_idx_0"]
-    assert [c.type.name for c in schema.columns] == ["bigint"]
-    assert meta.type == OutputType.index
-    assert meta.table_column_names == []
-    assert meta.table_index_column_names == ["_idx_0"]
-    assert meta.pd_column_level_names == []
-    assert meta.pd_index_level_names == [None]
-
-    with pytest.raises(AssertionError):
-        pandas_to_odps_schema(test_idx, unknown_as_string=True, ignore_index=True)
+    for ignore_idx in (False, True):
+        schema, meta = pandas_to_odps_schema(
+            test_idx, unknown_as_string=True, ignore_index=ignore_idx
+        )
+        assert [c.name for c in schema.columns] == ["_idx_0"]
+        assert [c.type.name for c in schema.columns] == ["bigint"]
+        assert meta.type == OutputType.index
+        assert meta.table_column_names == []
+        assert meta.table_index_column_names == ["_idx_0"]
+        assert meta.pd_column_level_names == []
+        assert meta.pd_index_level_names == [None]
 
     data = pd.MultiIndex.from_arrays(
         [np.random.choice(list("ABC"), 100), np.random.randint(0, 10, 100)],
@@ -177,6 +177,7 @@ def test_pandas_to_odps_schema_scalar(wrap_obj):
     test_scalar = _wrap_maxframe_obj(data, wrap=wrap_obj)
     if wrap_obj != "no":
         test_scalar.op.data = None
+
     schema, meta = pandas_to_odps_schema(test_scalar, unknown_as_string=True)
     assert schema.columns[0].name == "_idx_0"
     assert schema.columns[0].type.name == "double"
@@ -186,9 +187,6 @@ def test_pandas_to_odps_schema_scalar(wrap_obj):
     assert meta.pd_column_level_names == []
     assert meta.pd_index_level_names == [None]
 
-    with pytest.raises(AssertionError):
-        pandas_to_odps_schema(test_scalar, unknown_as_string=True, ignore_index=True)
-
 
 def test_odps_arrow_schema_conversion():
     odps_schema = odps_types.OdpsSchema(
@@ -211,10 +209,11 @@ def test_odps_arrow_schema_conversion():
             odps_types.Column("col16", "struct<a1: string, a2: map<string, bigint>>"),
             odps_types.Column("col17", "CHAR(15)"),
             odps_types.Column("col18", "VARCHAR(15)"),
+            odps_types.Column("col19", "decimal"),
         ]
     )
     arrow_schema = odps_schema_to_arrow_schema(odps_schema)
-    assert arrow_schema.names == [f"col{i}" for i in range(1, 19)]
+    assert arrow_schema.names == [f"col{i}" for i in range(1, 20)]
     assert arrow_schema.types == [
         pa.string(),
         pa.binary(),
@@ -234,6 +233,7 @@ def test_odps_arrow_schema_conversion():
         pa.struct([("a1", pa.string()), ("a2", pa.map_(pa.string(), pa.int64()))]),
         pa.string(),
         pa.string(),
+        pa.decimal128(38, 18),
     ]
 
     expected_odps_schema = odps_types.OdpsSchema(
@@ -256,6 +256,7 @@ def test_odps_arrow_schema_conversion():
             odps_types.Column("col16", "struct<a1: string, a2: map<string, bigint>>"),
             odps_types.Column("col17", "string"),
             odps_types.Column("col18", "string"),
+            odps_types.Column("col19", "decimal(38, 18)"),
         ]
     )