Python pandas.core.algorithms.take() Examples

def _pd_object_take(
        self,
        indices: Union[Sequence[int], np.ndarray],
        allow_fill: bool = False,
        fill_value: Optional[Any] = None,
    ) -> ExtensionArray:
        """Run take using object dtype and pandas' built-in algorithm.

        This is slow and should be avoided in future but is kept here as not all
        special cases are yet supported.
        """
        from pandas.core.algorithms import take

        data = self.astype(object)
        if allow_fill and fill_value is None:
            fill_value = self.dtype.na_value
        # fill value should always be translated from the scalar
        # type for the array, to the physical storage type for
        # the data, before passing to take.
        result = take(data, indices, fill_value=fill_value, allow_fill=allow_fill)
        return self._from_sequence(result, dtype=self.data.type) 

def test_2d_other_dtypes(self):
        arr = np.random.randn(10, 5).astype(np.float32)

        indexer = [1, 2, 3, -1]

        # axis=0
        result = algos.take_nd(arr, indexer, axis=0)
        expected = arr.take(indexer, axis=0)
        expected[-1] = np.nan
        tm.assert_almost_equal(result, expected)

        # axis=1
        result = algos.take_nd(arr, indexer, axis=1)
        expected = arr.take(indexer, axis=1)
        expected[:, -1] = np.nan
        tm.assert_almost_equal(result, expected) 

def _join_non_unique(self, other, how='left', return_indexers=False):
        from pandas.core.reshape.merge import _get_join_indexers

        left_idx, right_idx = _get_join_indexers([self._ndarray_values],
                                                 [other._ndarray_values],
                                                 how=how,
                                                 sort=True)

        left_idx = ensure_platform_int(left_idx)
        right_idx = ensure_platform_int(right_idx)

        join_index = np.asarray(self._ndarray_values.take(left_idx))
        mask = left_idx == -1
        np.putmask(join_index, mask, other._ndarray_values.take(right_idx))

        join_index = self._wrap_joined_index(join_index, other)

        if return_indexers:
            return join_index, left_idx, right_idx
        else:
            return join_index 

def _assert_take_fillable(self, values, indices, allow_fill=True,
                              fill_value=None, na_value=np.nan):
        """
        Internal method to handle NA filling of take.
        """
        indices = ensure_platform_int(indices)

        # only fill if we are passing a non-None fill_value
        if allow_fill and fill_value is not None:
            if (indices < -1).any():
                msg = ('When allow_fill=True and fill_value is not None, '
                       'all indices must be >= -1')
                raise ValueError(msg)
            taken = algos.take(values,
                               indices,
                               allow_fill=allow_fill,
                               fill_value=na_value)
        else:
            taken = values.take(indices)
        return taken 

def _join_non_unique(self, other, how='left', return_indexers=False):
        from pandas.core.reshape.merge import _get_join_indexers

        left_idx, right_idx = _get_join_indexers([self._ndarray_values],
                                                 [other._ndarray_values],
                                                 how=how,
                                                 sort=True)

        left_idx = _ensure_platform_int(left_idx)
        right_idx = _ensure_platform_int(right_idx)

        join_index = np.asarray(self._ndarray_values.take(left_idx))
        mask = left_idx == -1
        np.putmask(join_index, mask, other._ndarray_values.take(right_idx))

        join_index = self._wrap_joined_index(join_index, other)

        if return_indexers:
            return join_index, left_idx, right_idx
        else:
            return join_index 

def get_values(self):
        """
        Return the values.

        For internal compatibility with pandas formatting.

        Returns
        -------
        values : numpy array
            A numpy array of the same dtype as categorical.categories.dtype or
            Index if datetime / periods
        """
        # if we are a datetime and period index, return Index to keep metadata
        if is_datetimelike(self.categories):
            return self.categories.take(self._codes, fill_value=np.nan)
        elif is_integer_dtype(self.categories) and -1 in self._codes:
            return self.categories.astype("object").take(self._codes,
                                                         fill_value=np.nan)
        return np.array(self) 

def _concat_same_type(cls, to_concat):
        # type: (Sequence[ExtensionArray]) -> ExtensionArray
        """
        Concatenate multiple array

        Parameters
        ----------
        to_concat : sequence of this type

        Returns
        -------
        ExtensionArray
        """
        raise AbstractMethodError(cls)

    # The _can_hold_na attribute is set to True so that pandas internals
    # will use the ExtensionDtype.na_value as the NA value in operations
    # such as take(), reindex(), shift(), etc.  In addition, those results
    # will then be of the ExtensionArray subclass rather than an array
    # of objects 

def test_2d_other_dtypes(self):
        arr = np.random.randn(10, 5).astype(np.float32)

        indexer = [1, 2, 3, -1]

        # axis=0
        result = algos.take_nd(arr, indexer, axis=0)
        expected = arr.take(indexer, axis=0)
        expected[-1] = np.nan
        tm.assert_almost_equal(result, expected)

        # axis=1
        result = algos.take_nd(arr, indexer, axis=1)
        expected = arr.take(indexer, axis=1)
        expected[:, -1] = np.nan
        tm.assert_almost_equal(result, expected) 

def take(self, indices, axis=0, allow_fill=True,
             fill_value=None, **kwargs):
        if kwargs:
            nv.validate_take(tuple(), kwargs)
        indices = _ensure_platform_int(indices)
        if self._can_hold_na:
            taken = self._assert_take_fillable(self.values, indices,
                                               allow_fill=allow_fill,
                                               fill_value=fill_value,
                                               na_value=self._na_value)
        else:
            if allow_fill and fill_value is not None:
                msg = 'Unable to fill values because {0} cannot contain NA'
                raise ValueError(msg.format(self.__class__.__name__))
            taken = self.values.take(indices)
        return self._shallow_copy(taken) 

def _concat_same_type(cls, to_concat):
        # type: (Sequence[ExtensionArray]) -> ExtensionArray
        """
        Concatenate multiple array

        Parameters
        ----------
        to_concat : sequence of this type

        Returns
        -------
        ExtensionArray
        """
        raise AbstractMethodError(cls)

    # The _can_hold_na attribute is set to True so that pandas internals
    # will use the ExtensionDtype.na_value as the NA value in operations
    # such as take(), reindex(), shift(), etc.  In addition, those results
    # will then be of the ExtensionArray subclass rather than an array
    # of objects 

def get_values(self):
        """
        Return the values.

        For internal compatibility with pandas formatting.

        Returns
        -------
        values : numpy array
            A numpy array of the same dtype as categorical.categories.dtype or
            Index if datetime / periods
        """
        # if we are a datetime and period index, return Index to keep metadata
        if is_datetimelike(self.categories):
            return self.categories.take(self._codes, fill_value=np.nan)
        elif is_integer_dtype(self.categories) and -1 in self._codes:
            return self.categories.astype("object").take(self._codes,
                                                         fill_value=np.nan)
        return np.array(self) 

def _concat_same_type(cls, to_concat):
        # type: (Sequence[ExtensionArray]) -> ExtensionArray
        """Concatenate multiple array

        Parameters
        ----------
        to_concat : sequence of this type

        Returns
        -------
        ExtensionArray
        """
        raise AbstractMethodError(cls)

    # The _can_hold_na attribute is set to True so that pandas internals
    # will use the ExtensionDtype.na_value as the NA value in operations
    # such as take(), reindex(), shift(), etc.  In addition, those results
    # will then be of the ExtensionArray subclass rather than an array
    # of objects 

def test_2d_other_dtypes(self):
        arr = np.random.randn(10, 5).astype(np.float32)

        indexer = [1, 2, 3, -1]

        # axis=0
        result = algos.take_nd(arr, indexer, axis=0)
        expected = arr.take(indexer, axis=0)
        expected[-1] = np.nan
        tm.assert_almost_equal(result, expected)

        # axis=1
        result = algos.take_nd(arr, indexer, axis=1)
        expected = arr.take(indexer, axis=1)
        expected[:, -1] = np.nan
        tm.assert_almost_equal(result, expected) 

def _join_non_unique(self, other, how='left', return_indexers=False):
        from pandas.core.reshape.merge import _get_join_indexers

        left_idx, right_idx = _get_join_indexers([self._ndarray_values],
                                                 [other._ndarray_values],
                                                 how=how,
                                                 sort=True)

        left_idx = ensure_platform_int(left_idx)
        right_idx = ensure_platform_int(right_idx)

        join_index = np.asarray(self._ndarray_values.take(left_idx))
        mask = left_idx == -1
        np.putmask(join_index, mask, other._ndarray_values.take(right_idx))

        join_index = self._wrap_joined_index(join_index, other)

        if return_indexers:
            return join_index, left_idx, right_idx
        else:
            return join_index 

def _assert_take_fillable(self, values, indices, allow_fill=True,
                              fill_value=None, na_value=np.nan):
        """
        Internal method to handle NA filling of take.
        """
        indices = ensure_platform_int(indices)

        # only fill if we are passing a non-None fill_value
        if allow_fill and fill_value is not None:
            if (indices < -1).any():
                msg = ('When allow_fill=True and fill_value is not None, '
                       'all indices must be >= -1')
                raise ValueError(msg)
            taken = algos.take(values,
                               indices,
                               allow_fill=allow_fill,
                               fill_value=na_value)
        else:
            taken = values.take(indices)
        return taken 

def get_values(self):
        """ Return the values.

        For internal compatibility with pandas formatting.

        Returns
        -------
        values : numpy array
            A numpy array of the same dtype as categorical.categories.dtype or
            Index if datetime / periods
        """
        # if we are a datetime and period index, return Index to keep metadata
        if is_datetimelike(self.categories):
            return self.categories.take(self._codes, fill_value=np.nan)
        return np.array(self) 

def groupby(self, values):
        """
        Group the index labels by a given array of values.

        Parameters
        ----------
        values : array
            Values used to determine the groups.

        Returns
        -------
        groups : dict
            {group name -> group labels}
        """

        # TODO: if we are a MultiIndex, we can do better
        # that converting to tuples
        from .multi import MultiIndex
        if isinstance(values, MultiIndex):
            values = values.values
        values = _ensure_categorical(values)
        result = values._reverse_indexer()

        # map to the label
        result = {k: self.take(v) for k, v in compat.iteritems(result)}

        return result 

def _get_consensus_name(self, other):
        """
        Given 2 indexes, give a consensus name meaning
        we take the not None one, or None if the names differ.
        Return a new object if we are resetting the name
        """
        if self.name != other.name:
            if self.name is None or other.name is None:
                name = self.name or other.name
            else:
                name = None
            if self.name != name:
                return self._shallow_copy(name=name)
        return self 

def asof_locs(self, where, mask):
        """
        where : array of timestamps
        mask : array of booleans where data is not NA

        """
        locs = self.values[mask].searchsorted(where.values, side='right')

        locs = np.where(locs > 0, locs - 1, 0)
        result = np.arange(len(self))[mask].take(locs)

        first = mask.argmax()
        result[(locs == 0) & (where < self.values[first])] = -1

        return result 

def test_1d_with_out(self, dtype_can_hold_na, writeable):
        dtype, can_hold_na = dtype_can_hold_na

        data = np.random.randint(0, 2, 4).astype(dtype)
        data.flags.writeable = writeable

        indexer = [2, 1, 0, 1]
        out = np.empty(4, dtype=dtype)
        algos.take_1d(data, indexer, out=out)

        expected = data.take(indexer)
        tm.assert_almost_equal(out, expected)

        indexer = [2, 1, 0, -1]
        out = np.empty(4, dtype=dtype)

        if can_hold_na:
            algos.take_1d(data, indexer, out=out)
            expected = data.take(indexer)
            expected[3] = np.nan
            tm.assert_almost_equal(out, expected)
        else:
            with pytest.raises(TypeError, match=self.fill_error):
                algos.take_1d(data, indexer, out=out)

            # No Exception otherwise.
            data.take(indexer, out=out) 

def _from_factorized(cls, uniques, original):
        return original._constructor(original.categories.take(uniques),
                                     categories=original.categories,
                                     ordered=original.ordered) 

def test_2d_with_out(self, dtype_can_hold_na, writeable):
        dtype, can_hold_na = dtype_can_hold_na

        data = np.random.randint(0, 2, (5, 3)).astype(dtype)
        data.flags.writeable = writeable

        indexer = [2, 1, 0, 1]
        out0 = np.empty((4, 3), dtype=dtype)
        out1 = np.empty((5, 4), dtype=dtype)
        algos.take_nd(data, indexer, out=out0, axis=0)
        algos.take_nd(data, indexer, out=out1, axis=1)

        expected0 = data.take(indexer, axis=0)
        expected1 = data.take(indexer, axis=1)
        tm.assert_almost_equal(out0, expected0)
        tm.assert_almost_equal(out1, expected1)

        indexer = [2, 1, 0, -1]
        out0 = np.empty((4, 3), dtype=dtype)
        out1 = np.empty((5, 4), dtype=dtype)

        if can_hold_na:
            algos.take_nd(data, indexer, out=out0, axis=0)
            algos.take_nd(data, indexer, out=out1, axis=1)

            expected0 = data.take(indexer, axis=0)
            expected1 = data.take(indexer, axis=1)
            expected0[3, :] = np.nan
            expected1[:, 3] = np.nan

            tm.assert_almost_equal(out0, expected0)
            tm.assert_almost_equal(out1, expected1)
        else:
            for i, out in enumerate([out0, out1]):
                with pytest.raises(TypeError, match=self.fill_error):
                    algos.take_nd(data, indexer, out=out, axis=i)

                # No Exception otherwise.
                data.take(indexer, out=out, axis=i) 

def test_1d_other_dtypes(self):
        arr = np.random.randn(10).astype(np.float32)

        indexer = [1, 2, 3, -1]
        result = algos.take_1d(arr, indexer)
        expected = arr.take(indexer)
        expected[-1] = np.nan
        tm.assert_almost_equal(result, expected) 

def test_2d_with_out(self, dtype_can_hold_na, writeable):
        dtype, can_hold_na = dtype_can_hold_na

        data = np.random.randint(0, 2, (5, 3)).astype(dtype)
        data.flags.writeable = writeable

        indexer = [2, 1, 0, 1]
        out0 = np.empty((4, 3), dtype=dtype)
        out1 = np.empty((5, 4), dtype=dtype)
        algos.take_nd(data, indexer, out=out0, axis=0)
        algos.take_nd(data, indexer, out=out1, axis=1)

        expected0 = data.take(indexer, axis=0)
        expected1 = data.take(indexer, axis=1)
        tm.assert_almost_equal(out0, expected0)
        tm.assert_almost_equal(out1, expected1)

        indexer = [2, 1, 0, -1]
        out0 = np.empty((4, 3), dtype=dtype)
        out1 = np.empty((5, 4), dtype=dtype)

        if can_hold_na:
            algos.take_nd(data, indexer, out=out0, axis=0)
            algos.take_nd(data, indexer, out=out1, axis=1)

            expected0 = data.take(indexer, axis=0)
            expected1 = data.take(indexer, axis=1)
            expected0[3, :] = np.nan
            expected1[:, 3] = np.nan

            tm.assert_almost_equal(out0, expected0)
            tm.assert_almost_equal(out1, expected1)
        else:
            for i, out in enumerate([out0, out1]):
                with pytest.raises(TypeError, match=self.fill_error):
                    algos.take_nd(data, indexer, out=out, axis=i)

                # No Exception otherwise.
                data.take(indexer, out=out, axis=i) 

def test_take_coerces_list(self):
        arr = [1, 2, 3]
        result = algos.take(arr, [0, 0])
        expected = np.array([1, 1])
        tm.assert_numpy_array_equal(result, expected) 

def _from_factorized(cls, uniques, original):
        return original._constructor(original.categories.take(uniques),
                                     categories=original.categories,
                                     ordered=original.ordered) 

def test_take_coerces_list(self):
        arr = [1, 2, 3]
        result = algos.take(arr, [0, 0])
        expected = np.array([1, 1])
        tm.assert_numpy_array_equal(result, expected) 

def take(self, indices, allow_fill=False, fill_value=None):
        from pandas.core.algorithms import take

        result = take(self._ndarray, indices, allow_fill=allow_fill,
                      fill_value=fill_value)
        return type(self)(result) 

def remove_unused_categories(self, inplace=False):
        """
        Removes categories which are not used.

        Parameters
        ----------
        inplace : boolean (default: False)
           Whether or not to drop unused categories inplace or return a copy of
           this categorical with unused categories dropped.

        Returns
        -------
        cat : Categorical with unused categories dropped or None if inplace.

        See Also
        --------
        rename_categories
        reorder_categories
        add_categories
        remove_categories
        set_categories
        """
        inplace = validate_bool_kwarg(inplace, 'inplace')
        cat = self if inplace else self.copy()
        idx, inv = np.unique(cat._codes, return_inverse=True)

        if idx.size != 0 and idx[0] == -1:  # na sentinel
            idx, inv = idx[1:], inv - 1

        new_categories = cat.dtype.categories.take(idx)
        new_dtype = CategoricalDtype._from_fastpath(new_categories,
                                                    ordered=self.ordered)
        cat._dtype = new_dtype
        cat._codes = coerce_indexer_dtype(inv, new_dtype.categories)

        if not inplace:
            return cat 

def test_take_na_empty(self):
        result = algos.take(np.array([]), [-1, -1], allow_fill=True,
                            fill_value=0.0)
        expected = np.array([0., 0.])
        tm.assert_numpy_array_equal(result, expected)