odps/df/backends/frame.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 functools
import itertools
import operator

try:
    import numpy as np
    import pandas as pd

    has_pandas = True
except (ImportError, ValueError):
    has_pandas = False

from ...compat import u, six, izip as zip, Version
from ...config import options
from ...console import (
    get_console_size,
    in_interactive_session,
    in_ipython_frontend,
    in_qtconsole,
)
from ...models import TableSchema
from ...types import Partition
from ...tunnel.io.types import odps_type_to_arrow_type
from ...utils import to_str, to_text, deprecated
from . import formatter as fmt


class ResultFrame(six.Iterator):
    class ResultRecord(list):
        def __init__(self, columns, values):
            self._columns = columns
            self._column_id_by_name = dict([(v.name, k) for k, v in enumerate(columns)])
            super(ResultFrame.ResultRecord, self).__init__(values)

        def __getitem__(self, item):
            if isinstance(item, six.string_types):
                item = self._column_id_by_name[item]
            return list.__getitem__(self, item)

        def iteritems(self):
            for col, val in zip(self._columns, self):
                yield col.name, val

        def keys(self):
            return [c.name for c in self._columns]

    def __init__(self, data, columns=None, schema=None, index=None, pandas=True):
        if columns is None and schema is None:
            raise ValueError('Either columns or schema should be provided')

        if columns is None and schema is not None:
            columns = schema.columns

        self._columns = columns
        self._names = [to_text(c.name) for c in self._columns]
        self._types = [c.type for c in self._columns]
        self._index = index

        if has_pandas and pandas:
            if isinstance(data, pd.DataFrame):
                self._values = self._reset_pd_axes(data)
            else:
                self._values = pd.DataFrame(
                    [self._get_values(r) for r in data],
                    columns=self._names,
                    index=index,
                )
                self._cast_pd_types()
            self._index = self._values.index
            self._pandas = True
        else:
            if self._index is None:
                self._index = []
                self._values = []
                for i, r in zip(itertools.count(0), data):
                    self._values.append(self._get_values(r))
                    self._index.append(i)
            else:
                self._values = list(self._get_values(r) for r in data)
            self._pandas = False

        self._cursor = -1

    def _cast_pd_types(self):
        from .odpssql.types import df_type_to_odps_type
        from .pd.types import df_type_to_np_type

        assert isinstance(self._values, pd.DataFrame)

        if options.tunnel.pd_cast_mode == "arrow":
            dest_types = [
                pd.ArrowDtype(odps_type_to_arrow_type(df_type_to_odps_type(tp)))
                for tp in self._types
            ]
            self._values = self._values.astype(dict(zip(self._names, dest_types)))
        elif options.tunnel.pd_cast_mode == "numpy":
            dest_types = [df_type_to_np_type(tp) for tp in self._types]
            dest_df = self._values.copy()
            for (col_name, src_type), dest_type in zip(
                self._values.dtypes.items(), dest_types
            ):
                if src_type == np.dtype("O") and dest_type.kind in "fiu":
                    dest_type = np.dtype(float) if dest_type.kind in "iu" else dest_type
                    dest_df[col_name] = dest_df[col_name].astype(dest_type)
            self._values = dest_df

    def _reset_pd_axes(self, data):
        if Version(pd.__version__) < Version("0.23.0"):
            data = data.values
            return pd.DataFrame(
                [self._get_values(r) for r in data],
                columns=self._names,
                index=self._index,
            )

        ret_data = data
        if list(data.columns) != self._names:
            v = ret_data.set_axis(self._names, axis="columns")
            ret_data = v if v is not None else ret_data

        if self._index is not None and list(data.index) != list(self._index):
            if data is ret_data:
                ret_data = ret_data.set_axis(self._index, axis="index", inplace=False)
            else:
                ret_data.set_axis(self._index, axis="index", inplace=True)
        elif self._index is None and (
            not isinstance(data.index, pd.RangeIndex) or data.index.start != 0
        ):
            if data is ret_data:
                ret_data = ret_data.reset_index(drop=True, inplace=False)
            else:
                ret_data.reset_index(drop=True, inplace=True)
        return ret_data

    def _get_values(self, r):
        if hasattr(r, 'values'):
            return r.values
        return r

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

    @property
    def columns(self):
        return self._columns

    @property
    def names(self):
        return self._names

    @property
    def types(self):
        return self._types

    @property
    def schema(self):
        return TableSchema(
            columns=[col for col in self._columns if not isinstance(col, Partition)],
            partitions=[col for col in self._columns if isinstance(col, Partition)],
        )

    @property
    def index(self):
        return self._index

    @property
    def values(self):
        return self._values

    def get_column_data(self, item):
        item = to_text(item)
        if item not in self.names:
            return None
        if self._pandas:
            if item in self.values.index.names:
                return self.values.index.get_level_values(item)
            else:
                return self.values[item]
        else:
            col_id = list(idx for idx, c in self.names if c == item)[0]
            return [r[col_id] for r in self.values]

    if has_pandas:

        def to_pandas(self, wrap=False):
            from .. import DataFrame

            if wrap:
                return DataFrame(self.values)
            else:
                return self.values

    def __getattr__(self, item):
        col = self.get_column_data(item)
        if col is None:
            raise AttributeError("ResultFrame has no column '{0}'".format(item))
        return col

    def __getitem__(self, item):
        if isinstance(item, six.integer_types):
            if self._pandas:
                return self._values.iloc[item]
            else:
                return self.ResultRecord(self._columns, self._values[item])
        elif isinstance(item, slice):
            if self._pandas:
                return self._values.iloc[item]
            else:
                return ResultFrame(
                    self._values[item],
                    columns=self._columns,
                    index=self._index[item],
                    pandas=self._pandas,
                )
        elif isinstance(item, tuple) and len(item) == 2:
            if self._pandas:
                return self._values.iloc[item]
            else:
                if isinstance(item[1], slice):
                    frame = self[item[0]]
                    values = [r[item[1]] for r in frame._values]
                    return ResultFrame(
                        values,
                        columns=self._columns[item[1]],
                        index=frame._index,
                        pandas=self._pandas,
                    )
                else:
                    values = [r[item[1]] for r in self[item[0]]._values]
                    return values

    @deprecated("Direct call of next(ResultFrame) is deprecated")
    def __next__(self):
        self._cursor += 1
        try:
            if self._pandas:
                return self._values.iloc[self._cursor]
            else:
                return self.ResultRecord(self._columns, self._values[self._cursor])
        except IndexError:
            raise StopIteration

    def __iter__(self):
        cursor = 0
        while True:
            try:
                yield self[cursor]
                cursor += 1
            except IndexError:
                return

    def concat(self, *frames, **kwargs):
        if len(frames) == 2 and isinstance(frames[1], int):
            frames, axis = frames[:1], frames[1]
        else:
            axis = kwargs.pop("axis", 0)
        if kwargs:
            raise TypeError("Cannot accept arguments %s" % ",".join(kwargs))

        if len(frames) == 0:
            return self

        if self._pandas:
            try:
                from pandas import concat
            except ImportError:
                from pandas.tools.merge import concat

            values = (self._values,) + tuple(frame._values for frame in frames)
            pd_data = concat(values, axis=axis)
            if axis == 0:
                columns = self._columns
            else:
                columns = functools.reduce(
                    operator.add,
                    (self._columns,) + tuple(frame._columns for frame in frames),
                )
            return ResultFrame(pd_data, columns=columns, pandas=True)
        else:
            if axis == 0:
                if any(self._columns != frame._columns for frame in frames):
                    raise ValueError('Cannot concat two frame of different columns')

                values = functools.reduce(
                    operator.add,
                    (self._values,) + tuple(frame._values for frame in frames),
                )
                indices = functools.reduce(
                    operator.add,
                    (self._index,) + tuple(frame._index for frame in frames),
                )
                return ResultFrame(
                    values, columns=self._columns, index=indices, pandas=self._pandas
                )
            else:
                if any(self._index != frame._index for frame in frames):
                    raise ValueError('Cannot concat two frames of different indexes')

                sub_tuple = (self._values,) + tuple(frame._values for frame in frames)
                values = [
                    functools.reduce(operator.add, vals) for vals in zip(*sub_tuple)
                ]
                columns = functools.reduce(
                    operator.add,
                    (self._columns,) + tuple(frame._columns for frame in frames),
                )
                return ResultFrame(
                    values, columns, index=self._index, pandas=self._pandas
                )

    @property
    def dtypes(self):
        return [it.type for it in self._columns]

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

    def _repr_fits_vertical_(self):
        """
        Check length against max_rows.
        """
        max_rows = options.display.max_rows
        return len(self) <= max_rows

    def _repr_fits_horizontal_(self, ignore_width=False):
        """
        Check if full repr fits in horizontal boundaries imposed by the display
        options width and max_columns. In case off non-interactive session, no
        boundaries apply.

        ignore_width is here so ipnb+HTML output can behave the way
        users expect. display.max_columns remains in effect.
        GH3541, GH3573
        """

        width, height = get_console_size()
        max_columns = options.display.max_columns
        nb_columns = len(self.columns)

        # exceed max columns
        if (max_columns and nb_columns > max_columns) or (
            (not ignore_width) and width and nb_columns > (width // 2)
        ):
            return False

        if (
            ignore_width  # used by repr_html under IPython notebook
            # scripts ignore terminal dims
            or not in_interactive_session()
        ):
            return True

        if options.display.width is not None or in_ipython_frontend():
            # check at least the column row for excessive width
            max_rows = 1
        else:
            max_rows = options.display.max_rows

        # when auto-detecting, so width=None and not in ipython front end
        # check whether repr fits horizontal by actualy checking
        # the width of the rendered repr
        buf = six.StringIO()

        # only care about the stuff we'll actually print out
        # and to_string on entire frame may be expensive
        d = self

        if not (max_rows is None):  # unlimited rows
            # min of two, where one may be None
            d = d[: min(max_rows, len(d))]
        else:
            return True

        d.to_string(buf=buf)
        value = buf.getvalue()
        repr_width = max([len(l) for l in value.split('\n')])

        return repr_width < width

    def __repr__(self):
        return to_str(self.__unicode__())

    def __unicode__(self):
        """
        Return a string representation for a particular DataFrame
        """
        if self._pandas:
            return to_text(repr(self._values))

        buf = six.StringIO(u(""))

        max_rows = options.display.max_rows
        max_cols = options.display.max_columns
        show_dimensions = options.display.show_dimensions
        if options.display.expand_frame_repr:
            width, _ = get_console_size()
        else:
            width = None
        self.to_string(
            buf=buf,
            max_rows=max_rows,
            max_cols=max_cols,
            line_width=width,
            show_dimensions=show_dimensions,
        )

        return to_str(buf.getvalue())

    def _repr_html_(self):
        """
        Return a html representation for a particular DataFrame.
        Mainly for IPython notebook.
        """
        # qtconsole doesn't report its line width, and also
        # behaves badly when outputting an HTML table
        # that doesn't fit the window, so disable it.
        # XXX: In IPython 3.x and above, the Qt console will not attempt to
        # display HTML, so this check can be removed when support for IPython 2.x
        # is no longer needed.

        if self._pandas and options.display.notebook_widget:
            from .. import DataFrame
            from ..ui import show_df_widget

            show_df_widget(DataFrame(self._values, schema=self.schema))

        if self._pandas:
            return self._values._repr_html_()

        if in_qtconsole():
            # 'HTML output is disabled in QtConsole'
            return None

        if options.display.notebook_repr_html:
            max_rows = options.display.max_rows
            max_cols = options.display.max_columns
            show_dimensions = options.display.show_dimensions

            return self.to_html(
                max_rows=max_rows,
                max_cols=max_cols,
                show_dimensions=show_dimensions,
                notebook=True,
            )
        else:
            return None

    def to_string(
        self,
        buf=None,
        columns=None,
        col_space=None,
        header=True,
        index=True,
        na_rep='NaN',
        formatters=None,
        float_format=None,
        sparsify=None,
        index_names=True,
        justify=None,
        line_width=None,
        max_rows=None,
        max_cols=None,
        show_dimensions=False,
    ):
        """
        Render a DataFrame to a console-friendly tabular output.
        """

        formatter = fmt.ResultFrameFormatter(
            self,
            buf=buf,
            columns=columns,
            col_space=col_space,
            na_rep=na_rep,
            formatters=formatters,
            float_format=float_format,
            sparsify=sparsify,
            justify=justify,
            index_names=index_names,
            header=header,
            index=index,
            line_width=line_width,
            max_rows=max_rows,
            max_cols=max_cols,
            show_dimensions=show_dimensions,
        )
        formatter.to_string()

        if buf is None:
            result = formatter.buf.getvalue()
            return result

    def to_html(
        self,
        buf=None,
        columns=None,
        col_space=None,
        header=True,
        index=True,
        na_rep='NaN',
        formatters=None,
        float_format=None,
        sparsify=None,
        index_names=True,
        justify=None,
        bold_rows=True,
        classes=None,
        escape=True,
        max_rows=None,
        max_cols=None,
        show_dimensions=False,
        notebook=False,
    ):
        """
        Render a DataFrame as an HTML table.

        `to_html`-specific options:

        bold_rows : boolean, default True
            Make the row labels bold in the output
        classes : str or list or tuple, default None
            CSS class(es) to apply to the resulting html table
        escape : boolean, default True
            Convert the characters <, >, and & to HTML-safe sequences.=
        max_rows : int, optional
            Maximum number of rows to show before truncating. If None, show
            all.
        max_cols : int, optional
            Maximum number of columns to show before truncating. If None, show
            all.

        """

        formatter = fmt.ResultFrameFormatter(
            self,
            buf=buf,
            columns=columns,
            col_space=col_space,
            na_rep=na_rep,
            formatters=formatters,
            float_format=float_format,
            sparsify=sparsify,
            justify=justify,
            index_names=index_names,
            header=header,
            index=index,
            bold_rows=bold_rows,
            escape=escape,
            max_rows=max_rows,
            max_cols=max_cols,
            show_dimensions=show_dimensions,
        )
        formatter.to_html(classes=classes, notebook=notebook)

        if buf is None:
            return formatter.buf.getvalue()