maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/odpsio/schema.py

@@ -0,0 +1,364 @@
+# Copyright 1999-2024 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 string
+from collections import defaultdict
+from typing import Any, Dict, Tuple
+
+import pandas as pd
+import pyarrow as pa
+from odps import types as odps_types
+from pandas.api import types as pd_types
+
+from ..core import TILEABLE_TYPE, OutputType
+from ..protocol import DataFrameTableMeta
+from ..tensor.core import TENSOR_TYPE
+
+_TEMP_TABLE_PREFIX = "tmp_mf_"
+
+_arrow_to_odps_types = {
+    pa.string(): odps_types.string,
+    pa.binary(): odps_types.binary,
+    pa.int8(): odps_types.tinyint,
+    pa.int16(): odps_types.smallint,
+    pa.int32(): odps_types.int_,
+    pa.int64(): odps_types.bigint,
+    pa.bool_(): odps_types.boolean,
+    pa.float32(): odps_types.float_,
+    pa.float64(): odps_types.double,
+    pa.date32(): odps_types.date,
+    pa.timestamp("ms"): odps_types.datetime,
+    pa.timestamp("ns"): odps_types.timestamp,
+}
+
+_odps_type_to_arrow = {
+    odps_types.string: pa.string(),
+    odps_types.binary: pa.binary(),
+    odps_types.tinyint: pa.int8(),
+    odps_types.smallint: pa.int16(),
+    odps_types.int_: pa.int32(),
+    odps_types.bigint: pa.int64(),
+    odps_types.boolean: pa.bool_(),
+    odps_types.float_: pa.float32(),
+    odps_types.double: pa.float64(),
+    odps_types.date: pa.date32(),
+    odps_types.datetime: pa.timestamp("ms"),
+    odps_types.timestamp: pa.timestamp("ns"),
+}
+
+
+def arrow_type_to_odps_type(
+    arrow_type: pa.DataType, col_name: str, unknown_as_string: bool = False
+) -> odps_types.DataType:
+    if arrow_type in _arrow_to_odps_types:
+        return _arrow_to_odps_types[arrow_type]
+    elif isinstance(arrow_type, pa.ListType):
+        return odps_types.Array(
+            arrow_type_to_odps_type(arrow_type.value_type, col_name, unknown_as_string)
+        )
+    elif isinstance(arrow_type, pa.MapType):
+        return odps_types.Map(
+            arrow_type_to_odps_type(arrow_type.key_type, col_name, unknown_as_string),
+            arrow_type_to_odps_type(arrow_type.item_type, col_name, unknown_as_string),
+        )
+    elif isinstance(arrow_type, pa.StructType):
+        type_dict = {}
+        for idx in range(arrow_type.num_fields):
+            field = arrow_type[idx]
+            type_dict[field.name] = arrow_type_to_odps_type(
+                field.type, col_name, unknown_as_string
+            )
+        return odps_types.Struct(type_dict)
+    elif isinstance(arrow_type, pa.Decimal128Type):
+        return odps_types.Decimal(arrow_type.precision, arrow_type.scale)
+
+    if unknown_as_string:
+        return odps_types.string
+    else:
+        raise TypeError(
+            "Unknown type {}, column name is {},"
+            "specify `unknown_as_string=True` "
+            "or `as_type` to set column dtype".format(arrow_type, col_name)
+        )
+
+
+def arrow_schema_to_odps_schema(
+    arrow_schema: pa.Schema, unknown_as_string: bool = False
+) -> odps_types.OdpsSchema:
+    odps_cols = []
+    for col_name, col_type in zip(arrow_schema.names, arrow_schema.types):
+        col_odps_type = arrow_type_to_odps_type(
+            col_type, col_name, unknown_as_string=unknown_as_string
+        )
+        odps_cols.append(odps_types.Column(col_name, col_odps_type))
+
+    return odps_types.OdpsSchema(odps_cols)
+
+
+def odps_type_to_arrow_type(
+    odps_type: odps_types.DataType, col_name: str
+) -> pa.DataType:
+    if odps_type in _odps_type_to_arrow:
+        col_type = _odps_type_to_arrow[odps_type]
+    else:
+        if isinstance(odps_type, odps_types.Array):
+            col_type = pa.list_(odps_type_to_arrow_type(odps_type.value_type, col_name))
+        elif isinstance(odps_type, odps_types.Map):
+            col_type = pa.map_(
+                odps_type_to_arrow_type(odps_type.key_type, col_name),
+                odps_type_to_arrow_type(odps_type.value_type, col_name),
+            )
+        elif isinstance(odps_type, odps_types.Struct):
+            fields = [
+                (k, odps_type_to_arrow_type(v, col_name))
+                for k, v in odps_type.field_types.items()
+            ]
+            col_type = pa.struct(fields)
+        elif isinstance(odps_type, odps_types.Decimal):
+            col_type = pa.decimal128(
+                odps_type.precision or odps_types.Decimal._max_precision,
+                odps_type.scale or odps_types.Decimal._max_scale,
+            )
+        elif isinstance(odps_type, (odps_types.Varchar, odps_types.Char)):
+            col_type = pa.string()
+        else:
+            raise TypeError(
+                "Unsupported type {}, column name is {}".format(odps_type, col_name)
+            )
+    return col_type
+
+
+def odps_schema_to_arrow_schema(
+    odps_schema: odps_types.OdpsSchema, with_partitions: bool = False
+) -> pa.Schema:
+    arrow_schema = []
+    cols = odps_schema.columns if with_partitions else odps_schema.simple_columns
+    for col in cols:
+        col_name = col.name
+        col_type = odps_type_to_arrow_type(col.type, col_name)
+        arrow_schema.append(pa.field(col_name, col_type))
+
+    return pa.schema(arrow_schema)
+
+
+def odps_schema_to_pandas_dtypes(
+    odps_schema: odps_types.OdpsSchema, with_partitions: bool = False
+) -> pd.Series:
+    arrow_schema = odps_schema_to_arrow_schema(
+        odps_schema, with_partitions=with_partitions
+    )
+    return arrow_schema.empty_table().to_pandas().dtypes
+
+
+def _is_scalar_object(df_obj: Any) -> bool:
+    return (
+        isinstance(df_obj, TENSOR_TYPE) and df_obj.shape == ()
+    ) or pd_types.is_scalar(df_obj)
+
+
+def _scalar_as_index(df_obj: Any) -> pd.Index:
+    if isinstance(df_obj, TILEABLE_TYPE):
+        return pd.Index([], dtype=df_obj.dtype)
+    else:
+        return pd.Index([df_obj])[:0]
+
+
+def pandas_to_odps_schema(
+    df_obj: Any,
+    unknown_as_string: bool = False,
+    ignore_index=False,
+) -> Tuple[odps_types.OdpsSchema, DataFrameTableMeta]:
+    from .. import dataframe as md
+    from .arrow import pandas_to_arrow
+
+    if _is_scalar_object(df_obj):
+        empty_index = None
+    elif hasattr(df_obj, "index_value"):
+        empty_index = df_obj.index_value.to_pandas()[:0]
+    elif not isinstance(df_obj, pd.Index):
+        empty_index = df_obj.index[:0]
+    else:
+        empty_index = df_obj[:0]
+
+    if hasattr(df_obj, "columns_value"):
+        empty_columns = df_obj.dtypes.index
+    elif hasattr(df_obj, "columns"):
+        empty_columns = df_obj.columns
+    else:
+        empty_columns = None
+
+    if isinstance(df_obj, (md.DataFrame, pd.DataFrame)):
+        empty_df_obj = pd.DataFrame(
+            [], columns=empty_columns, index=empty_index
+        ).astype(df_obj.dtypes)
+    elif isinstance(df_obj, (md.Series, pd.Series)):
+        empty_df_obj = pd.Series([], name=df_obj.name, index=empty_index).astype(
+            df_obj.dtype
+        )
+    elif isinstance(df_obj, (md.Index, pd.Index)):
+        empty_df_obj = empty_index
+    else:
+        empty_df_obj = df_obj
+
+    arrow_data, table_meta = pandas_to_arrow(empty_df_obj, ignore_index=ignore_index)
+    return (
+        arrow_schema_to_odps_schema(
+            arrow_data.schema, unknown_as_string=unknown_as_string
+        ),
+        table_meta,
+    )
+
+
+def build_table_column_name(
+    col_idx: int, pd_col_name: Any, records: Dict[str, str]
+) -> str:
+    """
+    Convert column name to MaxCompute acceptable names
+
+    Parameters
+    ----------
+    col_idx:
+        index of the column
+    pd_col_name:
+        column name in pandas
+    records:
+        record for existing columns
+
+    Returns
+    -------
+        converted column name
+    """
+
+    def _is_col_name_legal(name: str):
+        if len(name) < 1 or len(name) > 128:
+            return False
+        if name[0] not in string.ascii_letters and name[0] != "_":
+            return False
+        for ch in name:
+            if ch not in string.digits and ch not in string.ascii_letters and ch != "_":
+                return False
+        return True
+
+    try:
+        return records[pd_col_name]
+    except KeyError:
+        pass
+
+    if isinstance(pd_col_name, str):
+        col_name = pd_col_name
+    elif isinstance(pd_col_name, tuple):
+        col_name = "_".join(str(x) for x in pd_col_name)
+    else:
+        col_name = str(pd_col_name)
+
+    col_name = col_name.lower()
+    if not _is_col_name_legal(col_name):
+        col_name = f"_column_{col_idx}"
+
+    records[pd_col_name] = col_name
+    return col_name
+
+
+def build_dataframe_table_meta(
+    df_obj: Any, ignore_index: bool = False
+) -> DataFrameTableMeta:
+    from .. import dataframe as md
+
+    col_to_count = defaultdict(lambda: 0)
+    col_to_idx = defaultdict(lambda: 0)
+    pd_col_to_col_name = dict()
+    if isinstance(df_obj, (md.DataFrame, pd.DataFrame)):
+        obj_type = OutputType.dataframe
+    elif isinstance(df_obj, (md.Series, pd.Series)):
+        obj_type = OutputType.series
+    elif isinstance(df_obj, (md.Index, pd.Index)):
+        obj_type = OutputType.index
+    elif _is_scalar_object(df_obj):
+        obj_type = OutputType.scalar
+    else:  # pragma: no cover
+        raise TypeError(f"Cannot accept type {type(df_obj)}")
+
+    assert not ignore_index or obj_type in (OutputType.dataframe, OutputType.series)
+
+    if obj_type == OutputType.scalar:
+        pd_dtypes = pd.Series([])
+        column_index_names = []
+        index_obj = _scalar_as_index(df_obj)
+    elif obj_type == OutputType.index:
+        pd_dtypes = pd.Series([])
+        column_index_names = []
+        index_obj = df_obj
+    elif obj_type == OutputType.series:
+        pd_dtypes = pd.Series([df_obj.dtype], index=[df_obj.name])
+        column_index_names = [None]
+        index_obj = df_obj.index
+    else:
+        pd_dtypes = df_obj.dtypes
+        column_index_names = list(pd_dtypes.index.names)
+        index_obj = df_obj.index
+
+    if isinstance(df_obj, TILEABLE_TYPE):
+        table_name = _TEMP_TABLE_PREFIX + str(df_obj.key)
+    else:
+        table_name = None
+
+    sql_columns = [None] * len(pd_dtypes)
+    pd_col_names = pd_dtypes.index
+    if obj_type == OutputType.series and df_obj.name is None:
+        # use special table column name for series
+        pd_col_names = ["_data"]
+    for idx, col in enumerate(pd_col_names):
+        sql_columns[idx] = col_name = build_table_column_name(
+            idx, col, pd_col_to_col_name
+        )
+        col_to_count[col_name] += 1
+
+    final_sql_columns = []
+    for col in sql_columns:
+        if col_to_count[col] > 1:
+            col_name = f"{col}_{col_to_idx[col]}"
+            col_to_idx[col] += 1
+            while col_name in col_to_count:
+                col_name = f"{col}_{col_to_idx[col]}"
+                col_to_idx[col] += 1
+            final_sql_columns.append(col_name)
+        else:
+            final_sql_columns.append(col)
+
+    if hasattr(index_obj, "index_value"):
+        pd_index_val = index_obj.index_value.to_pandas()
+    else:
+        pd_index_val = index_obj
+
+    if hasattr(pd_index_val, "dtypes"):
+        index_dtypes = pd.Series(pd_index_val.dtypes.values, index=pd_index_val.names)
+    else:
+        index_dtypes = pd.Series([pd_index_val.dtype], index=pd_index_val.names)
+
+    if ignore_index:
+        table_index_column_names = []
+        pd_index_dtypes = pd.Series([], index=[])
+    else:
+        table_index_column_names = [f"_idx_{i}" for i in range(len(index_obj.names))]
+        pd_index_dtypes = index_dtypes
+
+    return DataFrameTableMeta(
+        table_name=table_name,
+        type=obj_type,
+        table_column_names=final_sql_columns,
+        table_index_column_names=table_index_column_names,
+        pd_column_dtypes=pd_dtypes,
+        pd_column_level_names=column_index_names,
+        pd_index_dtypes=pd_index_dtypes,
+    )

maxframe/odpsio/tableio.py

@@ -0,0 +1,322 @@
+# Copyright 1999-2024 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 os
+import time
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import List, Optional, Union
+
+import pyarrow as pa
+from odps import ODPS
+from odps.apis.storage_api import (
+    StorageApiArrowClient,
+    TableBatchScanResponse,
+    TableBatchWriteResponse,
+)
+from odps.types import PartitionSpec
+
+from ..env import ODPS_STORAGE_API_ENDPOINT
+from .schema import odps_schema_to_arrow_schema
+
+PartitionsType = Union[List[str], str, None]
+
+_DEFAULT_ROW_BATCH_SIZE = 4096
+
+
+class MCTableIO(ABC):
+    def __init__(self, odps: ODPS):
+        self._odps = odps
+
+    @abstractmethod
+    def open_reader(
+        self,
+        full_table_name: str,
+        partitions: PartitionsType = None,
+        columns: Optional[List[str]] = None,
+        start: Optional[int] = None,
+        count: Optional[int] = None,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def open_writer(
+        self,
+        full_table_name: str,
+        partition: Optional[str] = None,
+        overwrite: bool = True,
+    ):
+        raise NotImplementedError
+
+
+class HaloTableArrowReader:
+    def __init__(
+        self,
+        client: StorageApiArrowClient,
+        scan_info: TableBatchScanResponse,
+        schema: pa.Schema,
+        start: Optional[int] = None,
+        count: Optional[int] = None,
+        row_batch_size: Optional[int] = None,
+    ):
+        self._client = client
+        self._scan_info = scan_info
+
+        self._cur_split_id = -1
+        self._cur_reader = None
+
+        self._schema = schema
+
+        self._start = start
+        self._count = count
+        self._cursor = 0
+        self._row_batch_size = row_batch_size
+
+    @property
+    def count(self) -> int:
+        return self._count
+
+    def _open_next_reader(self):
+        from odps.apis.storage_api import ReadRowsRequest
+
+        if 0 <= self._scan_info.split_count <= self._cur_split_id + 1:
+            # scan by split
+            self._cur_reader = None
+            return
+        elif self._count is not None and self._cursor >= self._count:
+            # scan by range
+            self._cur_reader = None
+            return
+
+        read_rows_kw = {}
+        if self._start is not None:
+            read_rows_kw["row_index"] = self._start + self._cursor
+            read_rows_kw["row_count"] = min(
+                self._row_batch_size, self._count - self._cursor
+            )
+            self._cursor = min(self._count, self._cursor + self._row_batch_size)
+
+        req = ReadRowsRequest(
+            session_id=self._scan_info.session_id,
+            split_index=self._cur_split_id + 1,
+            **read_rows_kw,
+        )
+        self._cur_reader = self._client.read_rows_arrow(req)
+        self._cur_split_id += 1
+
+    def read(self):
+        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:
+            batch = self._cur_reader.read()
+            if batch is not None:
+                return batch
+            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 HaloTableArrowWriter:
+    def __init__(
+        self, client: StorageApiArrowClient, write_info: TableBatchWriteResponse
+    ):
+        self._client = client
+        self._write_info = write_info
+
+        self._writer = None
+
+    def open(self):
+        from odps.apis.storage_api import WriteRowsRequest
+
+        self._writer = self._client.write_rows_arrow(
+            WriteRowsRequest(self._write_info.session_id)
+        )
+
+    def write(self, batch):
+        if isinstance(batch, pa.Table):
+            for b in batch.to_batches():
+                self._writer.write(b)
+        else:
+            self._writer.write(batch)
+
+    def close(self):
+        commit_msg, is_success = self._writer.finish()
+        if not is_success:
+            raise IOError(commit_msg)
+        return commit_msg
+
+
+class HaloTableIO(MCTableIO):
+    _storage_api_endpoint = os.getenv(ODPS_STORAGE_API_ENDPOINT)
+
+    @staticmethod
+    def _convert_partitions(partitions: PartitionsType) -> Optional[List[str]]:
+        if partitions is None:
+            return []
+        elif isinstance(partitions, (str, PartitionSpec)):
+            partitions = [partitions]
+        return [
+            "/".join(f"{k}={v}" for k, v in PartitionSpec(pt).items())
+            for pt in partitions
+        ]
+
+    def get_table_record_count(
+        self, full_table_name: str, partitions: PartitionsType = None
+    ):
+        from odps.apis.storage_api import SplitOptions, TableBatchScanRequest
+
+        table = self._odps.get_table(full_table_name)
+        client = StorageApiArrowClient(
+            self._odps, table, rest_endpoint=self._storage_api_endpoint
+        )
+
+        split_option = SplitOptions.SplitMode.SIZE
+
+        scan_kw = {
+            "required_partitions": self._convert_partitions(partitions),
+            "split_options": SplitOptions.get_default_options(split_option),
+        }
+
+        # todo add more options for partition column handling
+        req = TableBatchScanRequest(**scan_kw)
+        resp = client.create_read_session(req)
+        return resp.record_count
+
+    @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,
+    ):
+        from odps.apis.storage_api import (
+            SessionRequest,
+            SplitOptions,
+            Status,
+            TableBatchScanRequest,
+        )
+
+        table = self._odps.get_table(full_table_name)
+        client = StorageApiArrowClient(
+            self._odps, table, rest_endpoint=self._storage_api_endpoint
+        )
+
+        split_option = SplitOptions.SplitMode.SIZE
+        if start is not None or stop is not None:
+            split_option = SplitOptions.SplitMode.ROW_OFFSET
+
+        scan_kw = {
+            "required_partitions": self._convert_partitions(partitions),
+            "split_options": SplitOptions.get_default_options(split_option),
+        }
+        columns = columns or [c.name for c in table.table_schema.simple_columns]
+        scan_kw["required_data_columns"] = columns
+        if partition_columns is True:
+            scan_kw["required_partition_columns"] = [
+                c.name for c in table.table_schema.partitions
+            ]
+        else:
+            scan_kw["required_partition_columns"] = partition_columns
+
+        # todo add more options for partition column handling
+        req = TableBatchScanRequest(**scan_kw)
+        resp = client.create_read_session(req)
+
+        session_id = resp.session_id
+        status = resp.status
+        while status == Status.WAIT:
+            resp = client.get_read_session(SessionRequest(session_id))
+            status = resp.status
+            time.sleep(1.0)
+
+        assert status == Status.OK
+
+        count = None
+        if start is not None or stop is not None:
+            if reverse_range:
+                start = start if start is not None else resp.record_count - 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 resp.record_count
+            start = start if start >= 0 else resp.record_count + start
+            stop = stop if stop >= 0 else resp.record_count + stop
+            if reverse_range:
+                count = start - stop
+                start = stop + 1
+            else:
+                count = stop - start
+
+        yield HaloTableArrowReader(
+            client,
+            resp,
+            schema=odps_schema_to_arrow_schema(table.table_schema),
+            start=start,
+            count=count,
+            row_batch_size=row_batch_size,
+        )
+
+    @contextmanager
+    def open_writer(
+        self,
+        full_table_name: str,
+        partition: Optional[str] = None,
+        overwrite: bool = True,
+    ):
+        from odps.apis.storage_api import (
+            SessionRequest,
+            SessionStatus,
+            TableBatchWriteRequest,
+        )
+
+        table = self._odps.get_table(full_table_name)
+        client = StorageApiArrowClient(
+            self._odps, table, rest_endpoint=self._storage_api_endpoint
+        )
+
+        part_strs = self._convert_partitions(partition)
+        part_str = part_strs[0] if part_strs else None
+        req = TableBatchWriteRequest(partition_spec=part_str, overwrite=overwrite)
+        resp = client.create_write_session(req)
+
+        session_id = resp.session_id
+        writer = HaloTableArrowWriter(client, resp)
+        writer.open()
+
+        yield writer
+
+        commit_msg = writer.close()
+        resp = client.commit_write_session(
+            SessionRequest(session_id=session_id), [commit_msg]
+        )
+        while resp.session_status == SessionStatus.COMMITTING:
+            resp = client.get_write_session(SessionRequest(session_id=session_id))
+        assert resp.session_status == SessionStatus.COMMITTED

maxframe/odpsio/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 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.