indexing.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""
A loc indexer for pandas-on-Spark DataFrame/Series.
"""
from abc import ABCMeta, abstractmethod
from collections.abc import Iterable
from functools import reduce
from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union, cast, Sized

import pandas as pd
from pandas.api.types import is_list_like  # type: ignore[attr-defined]
from pyspark.sql import functions as F, Column
from pyspark.sql.types import BooleanType, LongType, DataType
from pyspark.errors import AnalysisException
import numpy as np

from pyspark import pandas as ps  # noqa: F401
from pyspark.pandas._typing import Label, Name, Scalar
from pyspark.pandas.internal import (
    DEFAULT_SERIES_NAME,
    InternalField,
    InternalFrame,
    NATURAL_ORDER_COLUMN_NAME,
    SPARK_DEFAULT_SERIES_NAME,
)
from pyspark.pandas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
from pyspark.pandas.utils import (
    is_name_like_tuple,
    is_name_like_value,
    lazy_property,
    name_like_string,
    same_anchor,
    scol_for,
    spark_column_equals,
    verify_temp_column_name,
)

if TYPE_CHECKING:
    from pyspark.pandas.frame import DataFrame
    from pyspark.pandas.generic import Frame
    from pyspark.pandas.series import Series


class IndexerLike:
    def __init__(self, psdf_or_psser: "Frame"):
        from pyspark.pandas.frame import DataFrame
        from pyspark.pandas.series import Series

        assert isinstance(
            psdf_or_psser, (DataFrame, Series)
        ), "unexpected argument type: {}".format(type(psdf_or_psser))
        self._psdf_or_psser = psdf_or_psser

    @property
    def _is_df(self) -> bool:
        from pyspark.pandas.frame import DataFrame

        return isinstance(self._psdf_or_psser, DataFrame)

    @property
    def _is_series(self) -> bool:
        from pyspark.pandas.series import Series

        return isinstance(self._psdf_or_psser, Series)

    @property
    def _psdf(self) -> "DataFrame":
        if self._is_df:
            return cast("DataFrame", self._psdf_or_psser)
        else:
            assert self._is_series
            return self._psdf_or_psser._psdf

    @property
    def _internal(self) -> InternalFrame:
        return self._psdf._internal


class AtIndexer(IndexerLike):
    """
    Access a single value for a row/column label pair.
    If the index is not unique, all matching pairs are returned as an array.
    Like ``loc``, in that both provide label-based lookups. Use ``at`` if you only need to
    get a single value in a DataFrame or Series.

    .. note:: Unlike pandas, pandas-on-Spark only allows using ``at`` to get values but not to
        set them.

    .. note:: Warning: If ``row_index`` matches a lot of rows, large amounts of data will be
        fetched, potentially causing your machine to run out of memory.

    Raises
    ------
    KeyError
        When label does not exist in DataFrame

    Examples
    --------
    >>> psdf = ps.DataFrame([[0, 2, 3], [0, 4, 1], [10, 20, 30]],
    ...                    index=[4, 5, 5], columns=['A', 'B', 'C'])
    >>> psdf
        A   B   C
    4   0   2   3
    5   0   4   1
    5  10  20  30

    Get value at specified row/column pair

    >>> psdf.at[4, 'B']
    2

    Get array if an index occurs multiple times

    >>> psdf.at[5, 'B']
    array([ 4, 20])
    """

    def __getitem__(self, key: Any) -> Union["Series", "DataFrame", Scalar]:
        if self._is_df:
            if not isinstance(key, tuple) or len(key) != 2:
                raise TypeError("Use DataFrame.at like .at[row_index, column_name]")
            row_sel, col_sel = key
        else:
            assert self._is_series, type(self._psdf_or_psser)
            if isinstance(key, tuple) and len(key) != 1:
                raise TypeError("Use Series.at like .at[row_index]")
            row_sel = key
            col_sel = self._psdf_or_psser._column_label

        if self._internal.index_level == 1:
            if not is_name_like_value(row_sel, allow_none=False, allow_tuple=False):
                raise ValueError("At based indexing on a single index can only have a single value")
            row_sel = (row_sel,)
        else:
            if not is_name_like_tuple(row_sel, allow_none=False):
                raise ValueError("At based indexing on multi-index can only have tuple values")

        if col_sel is not None:
            if not is_name_like_value(col_sel, allow_none=False):
                raise ValueError("At based indexing on multi-index can only have tuple values")
            if not is_name_like_tuple(col_sel):
                col_sel = (col_sel,)

        cond = reduce(
            lambda x, y: x & y,
            [scol == row for scol, row in zip(self._internal.index_spark_columns, row_sel)],
        )
        pdf = (
            self._internal.spark_frame.drop(NATURAL_ORDER_COLUMN_NAME)
            .filter(cond)
            .select(self._internal.spark_column_for(col_sel))
            .toPandas()
        )

        if len(pdf) < 1:
            raise KeyError(name_like_string(row_sel))

        values = pdf.iloc[:, 0].values
        return (
            values if (len(row_sel) < self._internal.index_level or len(values) > 1) else values[0]
        )


class iAtIndexer(IndexerLike):
    """
    Access a single value for a row/column pair by integer position.

    Like ``iloc``, in that both provide integer-based lookups. Use
    ``iat`` if you only need to get or set a single value in a DataFrame
    or Series.

    Raises
    ------
    KeyError
        When label does not exist in DataFrame

    Examples
    --------
    >>> df = ps.DataFrame([[0, 2, 3], [0, 4, 1], [10, 20, 30]],
    ...                   columns=['A', 'B', 'C'])
    >>> df
        A   B   C
    0   0   2   3
    1   0   4   1
    2  10  20  30

    Get value at specified row/column pair

    >>> df.iat[1, 2]
    1

    Get value within a series

    >>> psser = ps.Series([1, 2, 3], index=[10, 20, 30])
    >>> psser
    10    1
    20    2
    30    3
    dtype: int64

    >>> psser.iat[1]
    2
    """

    def __getitem__(self, key: Any) -> Union["Series", "DataFrame", Scalar]:
        if self._is_df:
            if not isinstance(key, tuple) or len(key) != 2:
                raise TypeError(
                    "Use DataFrame.iat like .iat[row_integer_position, column_integer_position]"
                )
            row_sel, col_sel = key
            if not isinstance(row_sel, int) or not isinstance(col_sel, int):
                raise ValueError("iAt based indexing can only have integer indexers")
            return self._psdf_or_psser.iloc[row_sel, col_sel]
        else:
            assert self._is_series, type(self._psdf_or_psser)
            if not isinstance(key, int) and len(key) != 1:
                raise TypeError("Use Series.iat like .iat[row_integer_position]")
            if not isinstance(key, int):
                raise ValueError("iAt based indexing can only have integer indexers")
            return self._psdf_or_psser.iloc[key]


class LocIndexerLike(IndexerLike, metaclass=ABCMeta):
    def _select_rows(self, rows_sel: Any) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """
        Dispatch the logic for select rows to more specific methods by `rows_sel` argument types.

        Parameters
        ----------
        rows_sel : the key specified to select rows.

        Returns
        -------
        Tuple of Spark column, int, int:

            * The Spark column for the condition to filter the rows.
            * The number of rows when the selection can be simplified by limit.
            * The remaining index rows if the result index size is shrunk.
        """
        from pyspark.pandas.series import Series

        if rows_sel is None:
            return None, None, None
        elif isinstance(rows_sel, Series):
            return self._select_rows_by_series(rows_sel)
        elif isinstance(rows_sel, Column):
            return self._select_rows_by_spark_column(rows_sel)
        elif isinstance(rows_sel, slice):
            if rows_sel == slice(None):
                # If slice is None - select everything, so nothing to do
                return None, None, None
            return self._select_rows_by_slice(rows_sel)
        elif isinstance(rows_sel, tuple):
            return self._select_rows_else(rows_sel)
        elif is_list_like(rows_sel):
            return self._select_rows_by_iterable(rows_sel)
        else:
            return self._select_rows_else(rows_sel)

    def _select_cols(
        self, cols_sel: Any, missing_keys: Optional[List[Name]] = None
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """
        Dispatch the logic for select columns to more specific methods by `cols_sel` argument types.

        Parameters
        ----------
        cols_sel : the key specified to select columns.

        Returns
        -------
        Tuple of list of column label, list of Spark columns, list of dtypes, bool:

            * The column labels selected.
            * The Spark columns selected.
            * The field metadata selected.
            * The boolean value whether Series should be returned or not.
            * The Series name if needed.
        """
        from pyspark.pandas.series import Series

        if cols_sel is None:
            column_labels = self._internal.column_labels
            data_spark_columns = self._internal.data_spark_columns
            data_fields = self._internal.data_fields
            return column_labels, data_spark_columns, data_fields, False, None
        elif isinstance(cols_sel, Series):
            return self._select_cols_by_series(cols_sel, missing_keys)
        elif isinstance(cols_sel, Column):
            return self._select_cols_by_spark_column(cols_sel, missing_keys)
        elif isinstance(cols_sel, slice):
            if cols_sel == slice(None):
                # If slice is None - select everything, so nothing to do
                column_labels = self._internal.column_labels
                data_spark_columns = self._internal.data_spark_columns
                data_fields = self._internal.data_fields
                return column_labels, data_spark_columns, data_fields, False, None
            return self._select_cols_by_slice(cols_sel, missing_keys)
        elif isinstance(cols_sel, tuple):
            return self._select_cols_else(cols_sel, missing_keys)
        elif is_list_like(cols_sel):
            return self._select_cols_by_iterable(cols_sel, missing_keys)
        else:
            return self._select_cols_else(cols_sel, missing_keys)

    # Methods for row selection

    @abstractmethod
    def _select_rows_by_series(
        self, rows_sel: "Series"
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """Select rows by `Series` type key."""
        pass

    @abstractmethod
    def _select_rows_by_spark_column(
        self, rows_sel: Column
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """Select rows by Spark `Column` type key."""
        pass

    @abstractmethod
    def _select_rows_by_slice(
        self, rows_sel: slice
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """Select rows by `slice` type key."""
        pass

    @abstractmethod
    def _select_rows_by_iterable(
        self, rows_sel: Iterable
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """Select rows by `Iterable` type key."""
        pass

    @abstractmethod
    def _select_rows_else(
        self, rows_sel: Any
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        """Select rows by other type key."""
        pass

    # Methods for col selection

    @abstractmethod
    def _select_cols_by_series(
        self, cols_sel: "Series", missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """Select columns by `Series` type key."""
        pass

    @abstractmethod
    def _select_cols_by_spark_column(
        self, cols_sel: Column, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """Select columns by Spark `Column` type key."""
        pass

    @abstractmethod
    def _select_cols_by_slice(
        self, cols_sel: slice, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """Select columns by `slice` type key."""
        pass

    @abstractmethod
    def _select_cols_by_iterable(
        self, cols_sel: Iterable, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """Select columns by `Iterable` type key."""
        pass

    @abstractmethod
    def _select_cols_else(
        self, cols_sel: Any, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        """Select columns by other type key."""
        pass

    def __getitem__(self, key: Any) -> Union["Series", "DataFrame"]:
        from pyspark.pandas.frame import DataFrame
        from pyspark.pandas.series import Series, first_series

        if self._is_series:
            if isinstance(key, Series) and not same_anchor(key, self._psdf_or_psser):
                name = self._psdf_or_psser.name or DEFAULT_SERIES_NAME
                psdf = self._psdf_or_psser.to_frame(name)
                temp_col = verify_temp_column_name(psdf, "__temp_col__")

                psdf[temp_col] = key
                return type(self)(psdf[name].rename(self._psdf_or_psser.name))[psdf[temp_col]]

            cond, limit, remaining_index = self._select_rows(key)
            if cond is None and limit is None:
                return self._psdf_or_psser

            column_label = self._psdf_or_psser._column_label
            column_labels = [column_label]
            data_spark_columns = [self._internal.spark_column_for(column_label)]
            data_fields = [self._internal.field_for(column_label)]
            returns_series = True
            series_name = self._psdf_or_psser.name
        else:
            assert self._is_df
            if isinstance(key, tuple):
                if len(key) != 2:
                    raise SparkPandasIndexingError("Only accepts pairs of candidates")
                rows_sel, cols_sel = key
            else:
                rows_sel = key
                cols_sel = None

            if isinstance(rows_sel, Series) and not same_anchor(rows_sel, self._psdf_or_psser):
                psdf = self._psdf_or_psser.copy()
                temp_col = verify_temp_column_name(cast("DataFrame", psdf), "__temp_col__")

                psdf[temp_col] = rows_sel
                return type(self)(psdf)[psdf[temp_col], cols_sel][list(self._psdf_or_psser.columns)]

            cond, limit, remaining_index = self._select_rows(rows_sel)
            (
                column_labels,
                data_spark_columns,
                data_fields,
                returns_series,
                series_name,
            ) = self._select_cols(cols_sel)

            if cond is None and limit is None and returns_series:
                psser = self._psdf_or_psser._psser_for(column_labels[0])
                if series_name is not None and series_name != psser.name:
                    psser = psser.rename(series_name)
                return psser

        if remaining_index is not None:
            index_spark_columns = self._internal.index_spark_columns[-remaining_index:]
            index_names = self._internal.index_names[-remaining_index:]
            index_fields = self._internal.index_fields[-remaining_index:]
        else:
            index_spark_columns = self._internal.index_spark_columns
            index_names = self._internal.index_names
            index_fields = self._internal.index_fields

        if len(column_labels) > 0:
            column_labels = column_labels.copy()
            column_labels_level = max(
                len(label) if label is not None else 1 for label in column_labels
            )
            none_column = 0
            for i, label in enumerate(column_labels):
                if label is None:
                    label = (none_column,)
                    none_column += 1
                if len(label) < column_labels_level:
                    label = tuple(list(label) + ([""]) * (column_labels_level - len(label)))
                column_labels[i] = label

            if i == 0 and none_column == 1:
                column_labels = [None]

            column_label_names = self._internal.column_label_names[-column_labels_level:]
        else:
            column_label_names = self._internal.column_label_names

        try:
            sdf = self._internal.spark_frame

            if cond is not None:
                index_columns = sdf.select(index_spark_columns).columns
                data_columns = sdf.select(data_spark_columns).columns
                sdf = sdf.filter(cond).select(index_spark_columns + data_spark_columns)
                index_spark_columns = [scol_for(sdf, col) for col in index_columns]
                data_spark_columns = [scol_for(sdf, col) for col in data_columns]

            if limit is not None:
                if limit >= 0:
                    sdf = sdf.limit(limit)
                else:
                    sdf = sdf.limit(sdf.count() + limit)
                sdf = sdf.drop(NATURAL_ORDER_COLUMN_NAME)
        except AnalysisException:
            raise KeyError(
                "[{}] don't exist in columns".format(
                    [col._jc.toString() for col in data_spark_columns]
                )
            )

        internal = InternalFrame(
            spark_frame=sdf,
            index_spark_columns=index_spark_columns,
            index_names=index_names,
            index_fields=index_fields,
            column_labels=column_labels,
            data_spark_columns=data_spark_columns,
            data_fields=data_fields,
            column_label_names=column_label_names,
        )
        psdf = DataFrame(internal)

        psdf_or_psser: Union[DataFrame, Series]
        if returns_series:
            psdf_or_psser = first_series(psdf)
            if series_name is not None and series_name != psdf_or_psser.name:
                psdf_or_psser = psdf_or_psser.rename(series_name)
        else:
            psdf_or_psser = psdf

        if remaining_index is not None and remaining_index == 0:
            pdf_or_pser = psdf_or_psser.head(2)._to_pandas()
            length = len(pdf_or_pser)
            if length == 0:
                raise KeyError(name_like_string(key))
            elif length == 1:
                return pdf_or_pser.iloc[0]
            else:
                return psdf_or_psser
        else:
            return psdf_or_psser

    def __setitem__(self, key: Any, value: Any) -> None:
        from pyspark.pandas.frame import DataFrame
        from pyspark.pandas.series import Series, first_series

        if self._is_series:
            if (
                isinstance(key, Series)
                and (isinstance(self, iLocIndexer) or not same_anchor(key, self._psdf_or_psser))
            ) or (
                isinstance(value, Series)
                and (isinstance(self, iLocIndexer) or not same_anchor(value, self._psdf_or_psser))
            ):
                if self._psdf_or_psser.name is None:
                    psdf = self._psdf_or_psser.to_frame()
                    column_label = psdf._internal.column_labels[0]
                else:
                    psdf = self._psdf_or_psser._psdf.copy()
                    column_label = self._psdf_or_psser._column_label
                temp_natural_order = verify_temp_column_name(psdf, "__temp_natural_order__")
                temp_key_col = verify_temp_column_name(psdf, "__temp_key_col__")
                temp_value_col = verify_temp_column_name(psdf, "__temp_value_col__")

                psdf[temp_natural_order] = F.monotonically_increasing_id()
                if isinstance(key, Series):
                    psdf[temp_key_col] = key
                if isinstance(value, Series):
                    psdf[temp_value_col] = value
                psdf = psdf.sort_values(temp_natural_order).drop(columns=temp_natural_order)

                psser = psdf._psser_for(column_label)
                if isinstance(key, Series):
                    key = F.col(
                        "`{}`".format(psdf[temp_key_col]._internal.data_spark_column_names[0])
                    )
                if isinstance(value, Series):
                    value = F.col(
                        "`{}`".format(psdf[temp_value_col]._internal.data_spark_column_names[0])
                    )

                type(self)(psser)[key] = value

                if self._psdf_or_psser.name is None:
                    psser = psser.rename()

                self._psdf_or_psser._psdf._update_internal_frame(
                    psser._psdf[
                        self._psdf_or_psser._psdf._internal.column_labels
                    ]._internal.resolved_copy,
                    check_same_anchor=False,
                )
                return

            if isinstance(value, DataFrame):
                raise ValueError("Incompatible indexer with DataFrame")

            cond, limit, remaining_index = self._select_rows(key)
            if cond is None:
                cond = F.lit(True)
            if limit is not None:
                cond = cond & (
                    self._internal.spark_frame[cast(iLocIndexer, self)._sequence_col] < F.lit(limit)
                )

            if isinstance(value, (Series, Column)):
                if remaining_index is not None and remaining_index == 0:
                    raise ValueError(
                        "No axis named {} for object type {}".format(key, type(value).__name__)
                    )
                if isinstance(value, Series):
                    value = value.spark.column
            else:
                value = F.lit(value)
            scol = (
                F.when(cond, value)
                .otherwise(self._internal.spark_column_for(self._psdf_or_psser._column_label))
                .alias(name_like_string(self._psdf_or_psser.name or SPARK_DEFAULT_SERIES_NAME))
            )

            internal = self._internal.with_new_spark_column(
                self._psdf_or_psser._column_label, scol  # TODO: dtype?
            )
            self._psdf_or_psser._psdf._update_internal_frame(internal, check_same_anchor=False)
        else:
            assert self._is_df

            if isinstance(key, tuple):
                if len(key) != 2:
                    raise SparkPandasIndexingError("Only accepts pairs of candidates")
                rows_sel, cols_sel = key
            else:
                rows_sel = key
                cols_sel = None

            if isinstance(value, DataFrame):
                if len(value.columns) == 1:
                    value = first_series(value)
                else:
                    raise ValueError("Only a dataframe with one column can be assigned")

            if (
                isinstance(rows_sel, Series)
                and (
                    isinstance(self, iLocIndexer) or not same_anchor(rows_sel, self._psdf_or_psser)
                )
            ) or (
                isinstance(value, Series)
                and (isinstance(self, iLocIndexer) or not same_anchor(value, self._psdf_or_psser))
            ):
                psdf = cast(DataFrame, self._psdf_or_psser.copy())
                temp_natural_order = verify_temp_column_name(psdf, "__temp_natural_order__")
                temp_key_col = verify_temp_column_name(psdf, "__temp_key_col__")
                temp_value_col = verify_temp_column_name(psdf, "__temp_value_col__")

                psdf[temp_natural_order] = F.monotonically_increasing_id()
                if isinstance(rows_sel, Series):
                    psdf[temp_key_col] = rows_sel
                if isinstance(value, Series):
                    psdf[temp_value_col] = value
                psdf = psdf.sort_values(temp_natural_order).drop(columns=temp_natural_order)

                if isinstance(rows_sel, Series):
                    rows_sel = F.col(
                        "`{}`".format(psdf[temp_key_col]._internal.data_spark_column_names[0])
                    )
                if isinstance(value, Series):
                    value = F.col(
                        "`{}`".format(psdf[temp_value_col]._internal.data_spark_column_names[0])
                    )

                type(self)(psdf)[rows_sel, cols_sel] = value

                self._psdf_or_psser._update_internal_frame(
                    psdf[list(self._psdf_or_psser.columns)]._internal.resolved_copy,
                    check_same_anchor=False,
                )
                return

            cond, limit, remaining_index = self._select_rows(rows_sel)
            missing_keys: List[Name] = []
            _, data_spark_columns, _, _, _ = self._select_cols(cols_sel, missing_keys=missing_keys)

            if cond is None:
                cond = F.lit(True)
            if limit is not None:
                cond = cond & (
                    self._internal.spark_frame[cast(iLocIndexer, self)._sequence_col] < F.lit(limit)
                )

            if isinstance(value, (Series, Column)):
                if remaining_index is not None and remaining_index == 0:
                    raise ValueError("Incompatible indexer with Series")
                if len(data_spark_columns) > 1:
                    raise ValueError("shape mismatch")
                if isinstance(value, Series):
                    value = value.spark.column
            else:
                value = F.lit(value)

            new_data_spark_columns = []
            new_fields = []
            for new_scol, spark_column_name, new_field in zip(
                self._internal.data_spark_columns,
                self._internal.data_spark_column_names,
                self._internal.data_fields,
            ):
                for scol in data_spark_columns:
                    if spark_column_equals(new_scol, scol):
                        new_scol = F.when(cond, value).otherwise(scol).alias(spark_column_name)
                        new_field = InternalField.from_struct_field(
                            self._internal.spark_frame.select(new_scol).schema[0],
                            use_extension_dtypes=new_field.is_extension_dtype,
                        )
                        break
                new_data_spark_columns.append(new_scol)
                new_fields.append(new_field)

            column_labels = self._internal.column_labels.copy()
            for missing in missing_keys:
                if is_name_like_tuple(missing):
                    label = cast(Label, missing)
                else:
                    label = cast(Label, (missing,))
                if len(label) < self._internal.column_labels_level:
                    label = tuple(
                        list(label) + ([""] * (self._internal.column_labels_level - len(label)))
                    )
                elif len(label) > self._internal.column_labels_level:
                    raise KeyError(
                        "Key length ({}) exceeds index depth ({})".format(
                            len(label), self._internal.column_labels_level
                        )
                    )
                column_labels.append(label)
                new_data_spark_columns.append(F.when(cond, value).alias(name_like_string(label)))
                new_fields.append(None)

            internal = self._internal.with_new_columns(
                new_data_spark_columns, column_labels=column_labels, data_fields=new_fields
            )
            self._psdf_or_psser._update_internal_frame(internal, check_same_anchor=False)


class LocIndexer(LocIndexerLike):
    """
    Access a group of rows and columns by label(s) or a boolean Series.

    ``.loc[]`` is primarily label based, but may also be used with a
    conditional boolean Series derived from the DataFrame or Series.

    Allowed inputs are:

    - A single label, e.g. ``5`` or ``'a'``, (note that ``5`` is
      interpreted as a *label* of the index, and **never** as an
      integer position along the index) for column selection.

    - A list or array of labels, e.g. ``['a', 'b', 'c']``.

    - A slice object with labels, e.g. ``'a':'f'``.

    - A conditional boolean Series derived from the DataFrame or Series

    - A boolean array of the same length as the column axis being sliced,
      e.g. ``[True, False, True]``.

    - An alignable boolean pandas Series to the column axis being sliced.
      The index of the key will be aligned before masking.

    Not allowed inputs which pandas allows are:

    - A boolean array of the same length as the row axis being sliced,
      e.g. ``[True, False, True]``.
    - A ``callable`` function with one argument (the calling Series, DataFrame
      or Panel) and that returns valid output for indexing (one of the above)

    .. note:: MultiIndex is not supported yet.

    .. note:: Note that contrary to usual python slices, **both** the
        start and the stop are included, and the step of the slice is not allowed.

    .. note:: With a list or array of labels for row selection,
        pandas-on-Spark behaves as a filter without reordering by the labels.

    See Also
    --------
    Series.loc : Access group of values using labels.

    Examples
    --------
    **Getting values**

    >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]],
    ...                   index=['cobra', 'viper', 'sidewinder'],
    ...                   columns=['max_speed', 'shield'])
    >>> df
                max_speed  shield
    cobra               1       2
    viper               4       5
    sidewinder          7       8

    Single label. Note this returns the row as a Series.

    >>> df.loc['viper']
    max_speed    4
    shield       5
    Name: viper, dtype: int64

    List of labels. Note using ``[[]]`` returns a DataFrame.
    Also note that pandas-on-Spark behaves just a filter without reordering by the labels.

    >>> df.loc[['viper', 'sidewinder']]
                max_speed  shield
    viper               4       5
    sidewinder          7       8

    >>> df.loc[['sidewinder', 'viper']]
                max_speed  shield
    viper               4       5
    sidewinder          7       8

    Single label for column.

    >>> df.loc['cobra', 'shield']
    2

    List of labels for row.

    >>> df.loc[['cobra'], 'shield']
    cobra    2
    Name: shield, dtype: int64

    List of labels for column.

    >>> df.loc['cobra', ['shield']]
    shield    2
    Name: cobra, dtype: int64

    List of labels for both row and column.

    >>> df.loc[['cobra'], ['shield']]
           shield
    cobra       2

    Slice with labels for row and single label for column.
    Note that both the start and stop of the slice are included.

    >>> df.loc['cobra':'viper', 'max_speed']
    cobra    1
    viper    4
    Name: max_speed, dtype: int64

    Conditional that returns a boolean Series

    >>> df.loc[df['shield'] > 6]
                max_speed  shield
    sidewinder          7       8

    Conditional that returns a boolean Series with column labels specified

    >>> df.loc[df['shield'] > 6, ['max_speed']]
                max_speed
    sidewinder          7

    A boolean array of the same length as the column axis being sliced.

    >>> df.loc[:, [False, True]]
                shield
    cobra            2
    viper            5
    sidewinder       8

    An alignable boolean Series to the column axis being sliced.

    >>> df.loc[:, pd.Series([False, True], index=['max_speed', 'shield'])]
                shield
    cobra            2
    viper            5
    sidewinder       8

    **Setting values**

    Setting value for all items matching the list of labels.

    >>> df.loc[['viper', 'sidewinder'], ['shield']] = 50
    >>> df
                max_speed  shield
    cobra               1       2
    viper               4      50
    sidewinder          7      50

    Setting value for an entire row

    >>> df.loc['cobra'] = 10
    >>> df
                max_speed  shield
    cobra              10      10
    viper               4      50
    sidewinder          7      50

    Set value for an entire column

    >>> df.loc[:, 'max_speed'] = 30
    >>> df
                max_speed  shield
    cobra              30      10
    viper              30      50
    sidewinder         30      50

    Set value for an entire list of columns

    >>> df.loc[:, ['max_speed', 'shield']] = 100
    >>> df
                max_speed  shield
    cobra             100     100
    viper             100     100
    sidewinder        100     100

    Set value with Series

    >>> df.loc[:, 'shield'] = df['shield'] * 2
    >>> df
                max_speed  shield
    cobra             100     200
    viper             100     200
    sidewinder        100     200

    **Getting values on a DataFrame with an index that has integer labels**

    Another example using integers for the index

    >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]],
    ...                   index=[7, 8, 9],
    ...                   columns=['max_speed', 'shield'])
    >>> df
       max_speed  shield
    7          1       2
    8          4       5
    9          7       8

    Slice with integer labels for rows. Note that both
    the start and stop of the slice are included.

    >>> df.loc[7:9]
       max_speed  shield
    7          1       2
    8          4       5
    9          7       8
    """

    @staticmethod
    def _NotImplemented(description: str) -> SparkPandasNotImplementedError:
        return SparkPandasNotImplementedError(
            description=description,
            pandas_function=".loc[..., ...]",
            spark_target_function="select, where",
        )

    def _select_rows_by_series(
        self, rows_sel: "Series"
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        assert isinstance(rows_sel.spark.data_type, BooleanType), rows_sel.spark.data_type
        return rows_sel.spark.column, None, None

    def _select_rows_by_spark_column(
        self, rows_sel: Column
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        spark_type = self._internal.spark_frame.select(rows_sel).schema[0].dataType
        assert isinstance(spark_type, BooleanType), spark_type
        return rows_sel, None, None

    def _select_rows_by_slice(
        self, rows_sel: slice
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        from pyspark.pandas.indexes import MultiIndex

        if rows_sel.step is not None:
            raise LocIndexer._NotImplemented("Cannot use step with Spark.")
        elif self._internal.index_level == 1:
            sdf = self._internal.spark_frame
            index = self._psdf_or_psser.index
            index_column = index.to_series()
            index_data_type = index_column.spark.data_type
            start = rows_sel.start
            stop = rows_sel.stop

            # get natural order from '__natural_order__' from start to stop
            # to keep natural order.
            start_and_stop = (
                sdf.select(index_column.spark.column, NATURAL_ORDER_COLUMN_NAME)
                .where(
                    (index_column.spark.column == F.lit(start).cast(index_data_type))
                    | (index_column.spark.column == F.lit(stop).cast(index_data_type))
                )
                .collect()
            )

            start = [row[1] for row in start_and_stop if row[0] == start]
            start = start[0] if len(start) > 0 else None

            stop = [row[1] for row in start_and_stop if row[0] == stop]
            stop = stop[-1] if len(stop) > 0 else None

            conds: List[Column] = []
            if start is not None:
                conds.append(F.col(NATURAL_ORDER_COLUMN_NAME) >= F.lit(start).cast(LongType()))
            if stop is not None:
                conds.append(F.col(NATURAL_ORDER_COLUMN_NAME) <= F.lit(stop).cast(LongType()))

            # if index order is not monotonic increasing or decreasing
            # and specified values don't exist in index, raise KeyError
            if (start is None and rows_sel.start is not None) or (
                stop is None and rows_sel.stop is not None
            ):

                inc = index_column.is_monotonic_increasing
                if inc is False:
                    dec = index_column.is_monotonic_decreasing

                if start is None and rows_sel.start is not None:
                    start = rows_sel.start
                    if inc is not False:
                        conds.append(
                            index_column.spark.column >= F.lit(start).cast(index_data_type)
                        )
                    elif dec is not False:
                        conds.append(
                            index_column.spark.column <= F.lit(start).cast(index_data_type)
                        )
                    else:
                        raise KeyError(rows_sel.start)
                if stop is None and rows_sel.stop is not None:
                    stop = rows_sel.stop
                    if inc is not False:
                        conds.append(index_column.spark.column <= F.lit(stop).cast(index_data_type))
                    elif dec is not False:
                        conds.append(index_column.spark.column >= F.lit(stop).cast(index_data_type))
                    else:
                        raise KeyError(rows_sel.stop)

            return reduce(lambda x, y: x & y, conds), None, None
        else:
            from pyspark.sql.types import StructType

            index = self._psdf_or_psser.index
            index_data_type = [  # type: ignore[assignment]
                f.dataType for f in cast(StructType, index.to_series().spark.data_type)
            ]

            start = rows_sel.start
            if start is not None:
                if not isinstance(start, tuple):
                    start = (start,)
                if len(start) == 0:
                    start = None
            stop = rows_sel.stop
            if stop is not None:
                if not isinstance(stop, tuple):
                    stop = (stop,)
                if len(stop) == 0:
                    stop = None

            depth = max(
                len(start) if start is not None else 0, len(stop) if stop is not None else 0
            )
            if depth == 0:
                return None, None, None
            elif (
                depth > self._internal.index_level
                or not index.droplevel(list(range(self._internal.index_level)[depth:])).is_monotonic
            ):
                raise KeyError(
                    "Key length ({}) was greater than MultiIndex sort depth".format(depth)
                )

            conds = []
            if start is not None:
                cond = F.lit(True)
                for scol, value, dt in list(
                    zip(
                        self._internal.index_spark_columns,
                        cast(Tuple[int, ...], start),
                        cast(List[DataType], index_data_type),
                    )
                )[::-1]:
                    compare = MultiIndex._comparator_for_monotonic_increasing(dt)
                    cond = F.when(scol.eqNullSafe(F.lit(value).cast(dt)), cond).otherwise(
                        compare(scol, F.lit(value).cast(dt), Column.__gt__)
                    )
                conds.append(cond)
            if stop is not None:
                cond = F.lit(True)
                for scol, value, dt in list(
                    zip(
                        self._internal.index_spark_columns,
                        cast(Tuple[int, ...], stop),
                        cast(List[DataType], index_data_type),
                    )
                )[::-1]:
                    compare = MultiIndex._comparator_for_monotonic_increasing(dt)
                    cond = F.when(scol.eqNullSafe(F.lit(value).cast(dt)), cond).otherwise(
                        compare(scol, F.lit(value).cast(dt), Column.__lt__)
                    )
                conds.append(cond)

            return reduce(lambda x, y: x & y, conds), None, None

    def _select_rows_by_iterable(
        self, rows_sel: Iterable
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        rows_sel = list(rows_sel)
        if len(rows_sel) == 0:
            return F.lit(False), None, None
        elif self._internal.index_level == 1:
            index_column = self._psdf_or_psser.index.to_series()
            index_data_type = index_column.spark.data_type
            if len(rows_sel) == 1:
                return (
                    index_column.spark.column == F.lit(rows_sel[0]).cast(index_data_type),
                    None,
                    None,
                )
            else:
                return (
                    index_column.spark.column.isin(
                        [F.lit(r).cast(index_data_type) for r in rows_sel]
                    ),
                    None,
                    None,
                )
        else:
            raise LocIndexer._NotImplemented("Cannot select with MultiIndex with Spark.")

    def _select_rows_else(
        self, rows_sel: Any
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        if not isinstance(rows_sel, tuple):
            rows_sel = (rows_sel,)
        if len(rows_sel) > self._internal.index_level:
            raise SparkPandasIndexingError("Too many indexers")

        rows = [scol == value for scol, value in zip(self._internal.index_spark_columns, rows_sel)]
        return (
            reduce(lambda x, y: x & y, rows),
            None,
            self._internal.index_level - len(rows_sel),
        )

    def _get_from_multiindex_column(
        self,
        key: Optional[Label],
        missing_keys: Optional[List[Name]],
        labels: Optional[List[Tuple[Label, Label]]] = None,
        recursed: int = 0,
    ) -> Tuple[List[Label], Optional[List[Column]], List[InternalField], bool, Optional[Name]]:
        """Select columns from multi-index columns."""
        assert isinstance(key, tuple)
        if labels is None:
            labels = [(label, label) for label in self._internal.column_labels]
        for k in key:
            labels = [
                (label, None if lbl is None else lbl[1:])
                for label, lbl in labels
                if (lbl is None and k is None) or (lbl is not None and lbl[0] == k)
            ]
            if len(labels) == 0:
                if missing_keys is None:
                    raise KeyError(k)
                else:
                    missing_keys.append(key)
                    return [], [], [], False, None

        if all(lbl is not None and len(lbl) > 0 and lbl[0] == "" for _, lbl in labels):
            # If the head is '', drill down recursively.
            labels = [(label, tuple([str(key), *lbl[1:]])) for i, (label, lbl) in enumerate(labels)]
            return self._get_from_multiindex_column((str(key),), missing_keys, labels, recursed + 1)
        else:
            returns_series = all(lbl is None or len(lbl) == 0 for _, lbl in labels)
            series_name: Optional[Name]
            if returns_series:
                label_set = set(label for label, _ in labels)
                assert len(label_set) == 1
                label = list(label_set)[0]
                column_labels = [label]
                data_spark_columns = [self._internal.spark_column_for(label)]
                data_fields = [self._internal.field_for(label)]
                if label is None:
                    series_name = None
                else:
                    if recursed > 0:
                        label = label[:-recursed]
                    series_name = label if len(label) > 1 else label[0]
            else:
                column_labels = [
                    None if lbl is None or lbl == (None,) else lbl for _, lbl in labels
                ]
                data_spark_columns = [self._internal.spark_column_for(label) for label, _ in labels]
                data_fields = [self._internal.field_for(label) for label, _ in labels]
                series_name = None

            return column_labels, data_spark_columns, data_fields, returns_series, series_name

    def _select_cols_by_series(
        self, cols_sel: "Series", missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        column_labels = cols_sel._internal.column_labels
        data_spark_columns = cols_sel._internal.data_spark_columns
        data_fields = cols_sel._internal.data_fields
        return column_labels, data_spark_columns, data_fields, True, None

    def _select_cols_by_spark_column(
        self, cols_sel: Column, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        column_labels: List[Label] = [(self._internal.spark_frame.select(cols_sel).columns[0],)]
        data_spark_columns = [cols_sel]
        return column_labels, data_spark_columns, None, True, None

    def _select_cols_by_slice(
        self, cols_sel: slice, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        start, stop = self._psdf_or_psser.columns.slice_locs(
            start=cols_sel.start, end=cols_sel.stop
        )
        column_labels = self._internal.column_labels[start:stop]
        data_spark_columns = self._internal.data_spark_columns[start:stop]
        data_fields = self._internal.data_fields[start:stop]
        return column_labels, data_spark_columns, data_fields, False, None

    def _select_cols_by_iterable(
        self, cols_sel: Iterable, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        from pyspark.pandas.series import Series

        if all(isinstance(key, Series) for key in cols_sel):
            column_labels = [key._column_label for key in cols_sel]
            data_spark_columns = [key.spark.column for key in cols_sel]
            data_fields = [key._internal.data_fields[0] for key in cols_sel]
        elif all(isinstance(key, Column) for key in cols_sel):
            column_labels = [
                (self._internal.spark_frame.select(col).columns[0],) for col in cols_sel
            ]
            data_spark_columns = list(cols_sel)
            data_fields = None
        elif all(isinstance(key, bool) for key in cols_sel) or all(
            isinstance(key, np.bool_) for key in cols_sel
        ):
            if len(cast(Sized, cols_sel)) != len(self._internal.column_labels):
                raise IndexError(
                    "Boolean index has wrong length: %s instead of %s"
                    % (len(cast(Sized, cols_sel)), len(self._internal.column_labels))
                )
            if isinstance(cols_sel, pd.Series):
                if not cols_sel.index.sort_values().equals(self._psdf.columns.sort_values()):
                    raise SparkPandasIndexingError(
                        "Unalignable boolean Series provided as indexer "
                        "(index of the boolean Series and of the indexed object do not match)"
                    )
                else:
                    column_labels = [
                        column_label
                        for column_label in self._internal.column_labels
                        if cols_sel[column_label if len(column_label) > 1 else column_label[0]]
                    ]
                    data_spark_columns = [
                        self._internal.spark_column_for(column_label)
                        for column_label in column_labels
                    ]
                    data_fields = [
                        self._internal.field_for(column_label) for column_label in column_labels
                    ]
            else:
                column_labels = [
                    self._internal.column_labels[i] for i, col in enumerate(cols_sel) if col
                ]
                data_spark_columns = [
                    self._internal.data_spark_columns[i] for i, col in enumerate(cols_sel) if col
                ]
                data_fields = [
                    self._internal.data_fields[i] for i, col in enumerate(cols_sel) if col
                ]
        elif any(isinstance(key, tuple) for key in cols_sel) and any(
            not is_name_like_tuple(key) for key in cols_sel
        ):
            raise TypeError(
                "Expected tuple, got {}".format(
                    type(set(key for key in cols_sel if not is_name_like_tuple(key)).pop())
                )
            )
        else:
            if missing_keys is None and all(isinstance(key, tuple) for key in cols_sel):
                level = self._internal.column_labels_level
                if any(len(key) != level for key in cols_sel):
                    raise ValueError("All the key level should be the same as column index level.")

            column_labels = []
            data_spark_columns = []
            data_fields = []
            for key in cols_sel:
                found = False
                for label in self._internal.column_labels:
                    if label == key or label[0] == key:
                        column_labels.append(label)
                        data_spark_columns.append(self._internal.spark_column_for(label))
                        data_fields.append(self._internal.field_for(label))
                        found = True
                if not found:
                    if missing_keys is None:
                        raise KeyError("['{}'] not in index".format(name_like_string(key)))
                    else:
                        missing_keys.append(key)

        return column_labels, data_spark_columns, data_fields, False, None

    def _select_cols_else(
        self, cols_sel: Any, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        if not is_name_like_tuple(cols_sel):
            cols_sel = (cols_sel,)
        return self._get_from_multiindex_column(cols_sel, missing_keys)


class iLocIndexer(LocIndexerLike):
    """
    Purely integer-location based indexing for selection by position.

    ``.iloc[]`` is primarily integer position based (from ``0`` to
    ``length-1`` of the axis), but may also be used with a conditional boolean Series.

    Allowed inputs are:

    - An integer for column selection, e.g. ``5``.
    - A list or array of integers for row selection with distinct index values,
      e.g. ``[3, 4, 0]``
    - A list or array of integers for column selection, e.g. ``[4, 3, 0]``.
    - A boolean array for column selection.
    - A slice object with ints for row and column selection, e.g. ``1:7``.

    Not allowed inputs which pandas allows are:

    - A list or array of integers for row selection with duplicated indexes,
      e.g. ``[4, 4, 0]``.
    - A boolean array for row selection.
    - A ``callable`` function with one argument (the calling Series, DataFrame
      or Panel) and that returns valid output for indexing (one of the above).
      This is useful in method chains when you don't have a reference to the
      calling object but would like to base your selection on some value.

    ``.iloc`` will raise ``IndexError`` if a requested indexer is
    out-of-bounds, except *slice* indexers which allow out-of-bounds
    indexing (this conforms with python/numpy *slice* semantics).

    See Also
    --------
    DataFrame.loc : Purely label-location based indexer for selection by label.
    Series.iloc : Purely integer-location based indexing for
                   selection by position.

    Examples
    --------

    >>> mydict = [{'a': 1, 'b': 2, 'c': 3, 'd': 4},
    ...           {'a': 100, 'b': 200, 'c': 300, 'd': 400},
    ...           {'a': 1000, 'b': 2000, 'c': 3000, 'd': 4000 }]
    >>> df = ps.DataFrame(mydict, columns=['a', 'b', 'c', 'd'])
    >>> df
          a     b     c     d
    0     1     2     3     4
    1   100   200   300   400
    2  1000  2000  3000  4000

    **Indexing just the rows**

    A scalar integer for row selection.

    >>> df.iloc[1]
    a    100
    b    200
    c    300
    d    400
    Name: 1, dtype: int64

    >>> df.iloc[[0]]
       a  b  c  d
    0  1  2  3  4

    With a `slice` object.

    >>> df.iloc[:3]
          a     b     c     d
    0     1     2     3     4
    1   100   200   300   400
    2  1000  2000  3000  4000

    **Indexing both axes**

    You can mix the indexer types for the index and columns. Use ``:`` to
    select the entire axis.

    With scalar integers.

    >>> df.iloc[:1, 1]
    0    2
    Name: b, dtype: int64

    With lists of integers.

    >>> df.iloc[:2, [1, 3]]
         b    d
    0    2    4
    1  200  400

    With `slice` objects.

    >>> df.iloc[:2, 0:3]
         a    b    c
    0    1    2    3
    1  100  200  300

    With a boolean array whose length matches the columns.

    >>> df.iloc[:, [True, False, True, False]]
          a     c
    0     1     3
    1   100   300
    2  1000  3000

    **Setting values**

    Setting value for all items matching the list of labels.

    >>> df.iloc[[1, 2], [1]] = 50
    >>> df
          a   b     c     d
    0     1   2     3     4
    1   100  50   300   400
    2  1000  50  3000  4000

    Setting value for an entire row

    >>> df.iloc[0] = 10
    >>> df
          a   b     c     d
    0    10  10    10    10
    1   100  50   300   400
    2  1000  50  3000  4000

    Set value for an entire column

    >>> df.iloc[:, 2] = 30
    >>> df
          a   b   c     d
    0    10  10  30    10
    1   100  50  30   400
    2  1000  50  30  4000

    Set value for an entire list of columns

    >>> df.iloc[:, [2, 3]] = 100
    >>> df
          a   b    c    d
    0    10  10  100  100
    1   100  50  100  100
    2  1000  50  100  100

    Set value with Series

    >>> df.iloc[:, 3] = df.iloc[:, 3] * 2
    >>> df
          a   b    c    d
    0    10  10  100  200
    1   100  50  100  200
    2  1000  50  100  200
    """

    @staticmethod
    def _NotImplemented(description: str) -> SparkPandasNotImplementedError:
        return SparkPandasNotImplementedError(
            description=description,
            pandas_function=".iloc[..., ...]",
            spark_target_function="select, where",
        )

    @lazy_property
    def _internal(self) -> "InternalFrame":
        # Use resolved_copy to fix the natural order.
        internal = super()._internal.resolved_copy
        sdf = InternalFrame.attach_distributed_sequence_column(
            internal.spark_frame, column_name=self._sequence_col
        )
        return internal.with_new_sdf(spark_frame=sdf.orderBy(NATURAL_ORDER_COLUMN_NAME))

    @lazy_property
    def _sequence_col(self) -> str:
        # Use resolved_copy to fix the natural order.
        internal = super()._internal.resolved_copy
        return verify_temp_column_name(internal.spark_frame, "__distributed_sequence_column__")

    def _select_rows_by_series(
        self, rows_sel: "Series"
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        raise iLocIndexer._NotImplemented(
            ".iloc requires numeric slice, conditional "
            "boolean Index or a sequence of positions as int, "
            "got {}".format(type(rows_sel))
        )

    def _select_rows_by_spark_column(
        self, rows_sel: Column
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        raise iLocIndexer._NotImplemented(
            ".iloc requires numeric slice, conditional "
            "boolean Index or a sequence of positions as int, "
            "got {}".format(type(rows_sel))
        )

    def _select_rows_by_slice(
        self, rows_sel: slice
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        def verify_type(i: int) -> None:
            if not isinstance(i, int):
                raise TypeError(
                    "cannot do slice indexing with these indexers [{}] of {}".format(i, type(i))
                )

        has_negative = False
        start = rows_sel.start
        if start is not None:
            verify_type(start)
            if start == 0:
                start = None
            elif start < 0:
                has_negative = True
        stop = rows_sel.stop
        if stop is not None:
            verify_type(stop)
            if stop < 0:
                has_negative = True

        step = rows_sel.step
        if step is not None:
            verify_type(step)
            if step == 0:
                raise ValueError("slice step cannot be zero")
        else:
            step = 1

        if start is None and step == 1:
            return None, stop, None

        sdf = self._internal.spark_frame
        sequence_scol = sdf[self._sequence_col]

        if has_negative or (step < 0 and start is None):
            cnt = sdf.count()

        cond = []
        if start is not None:
            if start < 0:
                start = start + cnt
            if step >= 0:
                cond.append(sequence_scol >= F.lit(start).cast(LongType()))
            else:
                cond.append(sequence_scol <= F.lit(start).cast(LongType()))
        if stop is not None:
            if stop < 0:
                stop = stop + cnt
            if step >= 0:
                cond.append(sequence_scol < F.lit(stop).cast(LongType()))
            else:
                cond.append(sequence_scol > F.lit(stop).cast(LongType()))
        if step != 1:
            if step > 0:
                start = start or 0
            else:
                start = start or (cnt - 1)
            cond.append(((sequence_scol - start) % F.lit(step).cast(LongType())) == F.lit(0))

        return reduce(lambda x, y: x & y, cond), None, None

    def _select_rows_by_iterable(
        self, rows_sel: Iterable
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        sdf = self._internal.spark_frame

        if any(isinstance(key, (int, np.int64, np.int32)) and key < 0 for key in rows_sel):
            offset = sdf.count()
        else:
            offset = 0

        new_rows_sel = []
        for key in list(rows_sel):
            if not isinstance(key, (int, np.int64, np.int32)):
                raise TypeError(
                    "cannot do positional indexing with these indexers [{}] of {}".format(
                        key, type(key)
                    )
                )
            if key < 0:
                key = key + offset
            new_rows_sel.append(key)

        if len(new_rows_sel) != len(set(new_rows_sel)):
            raise NotImplementedError(
                "Duplicated row selection is not currently supported; "
                "however, normalized index was [%s]" % new_rows_sel
            )

        if len(new_rows_sel) == 0:
            cond = F.lit(False)
        else:
            cond = sdf[self._sequence_col].isin(
                [F.lit(int(key)).cast(LongType()) for key in new_rows_sel]
            )
        return cond, None, None

    def _select_rows_else(
        self, rows_sel: Any
    ) -> Tuple[Optional[Column], Optional[int], Optional[int]]:
        if isinstance(rows_sel, int):
            sdf = self._internal.spark_frame
            return (sdf[self._sequence_col] == rows_sel), None, 0
        elif isinstance(rows_sel, tuple):
            raise SparkPandasIndexingError("Too many indexers")
        else:
            raise iLocIndexer._NotImplemented(
                ".iloc requires numeric slice, conditional "
                "boolean Index or a sequence of positions as int, "
                "got {}".format(type(rows_sel))
            )

    def _select_cols_by_series(
        self, cols_sel: "Series", missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        raise ValueError(
            "Location based indexing can only have [integer, integer slice, "
            "listlike of integers, boolean array] types, got {}".format(cols_sel)
        )

    def _select_cols_by_spark_column(
        self, cols_sel: Column, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        raise ValueError(
            "Location based indexing can only have [integer, integer slice, "
            "listlike of integers, boolean array] types, got {}".format(cols_sel)
        )

    def _select_cols_by_slice(
        self, cols_sel: slice, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        if all(
            s is None or isinstance(s, int) for s in (cols_sel.start, cols_sel.stop, cols_sel.step)
        ):
            column_labels = self._internal.column_labels[cols_sel]
            data_spark_columns = self._internal.data_spark_columns[cols_sel]
            data_fields = self._internal.data_fields[cols_sel]
            return column_labels, data_spark_columns, data_fields, False, None
        else:
            not_none = (
                cols_sel.start
                if cols_sel.start is not None
                else cols_sel.stop
                if cols_sel.stop is not None
                else cols_sel.step
            )
            raise TypeError(
                "cannot do slice indexing with these indexers {} of {}".format(
                    not_none, type(not_none)
                )
            )

    def _select_cols_by_iterable(
        self, cols_sel: Iterable, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        if all(isinstance(s, bool) for s in cols_sel):
            cols_sel = [i for i, s in enumerate(cols_sel) if s]
        if all(isinstance(s, int) for s in cols_sel):
            column_labels = [self._internal.column_labels[s] for s in cols_sel]
            data_spark_columns = [self._internal.data_spark_columns[s] for s in cols_sel]
            data_fields = [self._internal.data_fields[s] for s in cols_sel]
            return column_labels, data_spark_columns, data_fields, False, None
        else:
            raise TypeError("cannot perform reduce with flexible type")

    def _select_cols_else(
        self, cols_sel: Any, missing_keys: Optional[List[Name]]
    ) -> Tuple[
        List[Label],
        Optional[List[Column]],
        Optional[List[InternalField]],
        bool,
        Optional[Name],
    ]:
        if isinstance(cols_sel, int):
            if cols_sel > len(self._internal.column_labels):
                raise KeyError(cols_sel)
            column_labels = [self._internal.column_labels[cols_sel]]
            data_spark_columns = [self._internal.data_spark_columns[cols_sel]]
            data_fields = [self._internal.data_fields[cols_sel]]
            return column_labels, data_spark_columns, data_fields, True, None
        else:
            raise ValueError(
                "Location based indexing can only have [integer, integer slice, "
                "listlike of integers, boolean array] types, got {}".format(cols_sel)
            )

    def __setitem__(self, key: Any, value: Any) -> None:
        if not isinstance(value, Column) and is_list_like(value):
            iloc_item = self[key]
            if not is_list_like(key) or not is_list_like(iloc_item):
                raise ValueError("setting an array element with a sequence.")
            else:
                shape_iloc_item = iloc_item.shape
                len_iloc_item = shape_iloc_item[0]
                len_value = len(value)
                if len_iloc_item != len_value:
                    if self._is_series:
                        raise ValueError(
                            "cannot set using a list-like indexer with a different length than "
                            "the value"
                        )
                    else:
                        raise ValueError(
                            "shape mismatch: value array of shape ({},) could not be broadcast "
                            "to indexing result of shape {}".format(len_value, shape_iloc_item)
                        )
        super().__setitem__(key, value)
        # Update again with resolved_copy to drop extra columns.
        self._psdf._update_internal_frame(
            self._psdf._internal.resolved_copy, check_same_anchor=False
        )

        # Clean up implicitly cached properties to be able to reuse the indexer.
        del self._internal
        del self._sequence_col


def _test() -> None:
    import os
    import doctest
    import sys
    from pyspark.sql import SparkSession
    import pyspark.pandas.indexing

    os.chdir(os.environ["SPARK_HOME"])

    globs = pyspark.pandas.indexing.__dict__.copy()
    globs["ps"] = pyspark.pandas
    spark = (
        SparkSession.builder.master("local[4]")
        .appName("pyspark.pandas.indexing tests")
        .getOrCreate()
    )
    (failure_count, test_count) = doctest.testmod(
        pyspark.pandas.indexing,
        globs=globs,
        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
    )
    spark.stop()
    if failure_count:
        sys.exit(-1)


if __name__ == "__main__":
    _test()