def __getitem__(self, item): if isinstance(item, tuple): item = unpack_tuple_and_ellipses(item) if isinstance(item, numbers.Integral): return self.data[item] elif isinstance(item, slice) and item == slice(None): # Make sure we get a view return type(self)(self.data) elif isinstance(item, slice): # slice return type(self)(self.data[item]) elif not is_list_like(item): # e.g. "foo" or 2.5 # exception message copied from numpy raise IndexError( r"only integers, slices (`:`), ellipsis (`...`), numpy.newaxis " r"(`None`) and integer or boolean arrays are valid indices") else: item = pd.api.indexers.check_array_indexer(self, item) if is_bool_dtype(item.dtype): return self._from_sequence( [x for x, m in zip(self, item) if m]) # integer return type(self)([self.data[i] for i in item])
def __getitem__( self: ArrowStringArray, item: PositionalIndexer ) -> ArrowStringArray | ArrowStringScalarOrNAT: """Select a subset of self. Parameters ---------- item : int, slice, or ndarray * int: The position in 'self' to get. * slice: A slice object, where 'start', 'stop', and 'step' are integers or None * ndarray: A 1-d boolean NumPy ndarray the same length as 'self' Returns ------- item : scalar or ExtensionArray Notes ----- For scalar ``item``, return a scalar value suitable for the array's type. This should be an instance of ``self.dtype.type``. For slice ``key``, return an instance of ``ExtensionArray``, even if the slice is length 0 or 1. For a boolean mask, return an instance of ``ExtensionArray``, filtered to the values where ``item`` is True. """ item = check_array_indexer(self, item) if isinstance(item, np.ndarray): if not len(item): return type(self)(pa.chunked_array([], type=pa.string())) elif is_integer_dtype(item.dtype): return self.take(item) elif is_bool_dtype(item.dtype): return type(self)(self._data.filter(item)) else: raise IndexError( "Only integers, slices and integer or " "boolean arrays are valid indices." ) elif isinstance(item, tuple): item = unpack_tuple_and_ellipses(item) if is_scalar(item) and not is_integer(item): # e.g. "foo" or 2.5 # exception message copied from numpy raise IndexError( r"only integers, slices (`:`), ellipsis (`...`), numpy.newaxis " r"(`None`) and integer or boolean arrays are valid indices" ) # We are not an array indexer, so maybe e.g. a slice or integer # indexer. We dispatch to pyarrow. value = self._data[item] if isinstance(value, pa.ChunkedArray): return type(self)(value) else: return self._as_pandas_scalar(value)
def __getitem__(self, item: PositionalIndexer): """Select a subset of self. Parameters ---------- item : int, slice, or ndarray * int: The position in 'self' to get. * slice: A slice object, where 'start', 'stop', and 'step' are integers or None * ndarray: A 1-d boolean NumPy ndarray the same length as 'self' Returns ------- item : scalar or ExtensionArray Notes ----- For scalar ``item``, return a scalar value suitable for the array's type. This should be an instance of ``self.dtype.type``. For slice ``key``, return an instance of ``ExtensionArray``, even if the slice is length 0 or 1. For a boolean mask, return an instance of ``ExtensionArray``, filtered to the values where ``item`` is True. """ item = check_array_indexer(self, item) if isinstance(item, np.ndarray): if not len(item): # Removable once we migrate StringDtype[pyarrow] to ArrowDtype[string] if self._dtype.name == "string" and self._dtype.storage == "pyarrow": pa_dtype = pa.string() else: pa_dtype = self._dtype.pyarrow_dtype return type(self)(pa.chunked_array([], type=pa_dtype)) elif is_integer_dtype(item.dtype): return self.take(item) elif is_bool_dtype(item.dtype): return type(self)(self._data.filter(item)) else: raise IndexError("Only integers, slices and integer or " "boolean arrays are valid indices.") elif isinstance(item, tuple): item = unpack_tuple_and_ellipses(item) # error: Non-overlapping identity check (left operand type: # "Union[Union[int, integer[Any]], Union[slice, List[int], # ndarray[Any, Any]]]", right operand type: "ellipsis") if item is Ellipsis: # type: ignore[comparison-overlap] # TODO: should be handled by pyarrow? item = slice(None) if is_scalar(item) and not is_integer(item): # e.g. "foo" or 2.5 # exception message copied from numpy raise IndexError( r"only integers, slices (`:`), ellipsis (`...`), numpy.newaxis " r"(`None`) and integer or boolean arrays are valid indices") # We are not an array indexer, so maybe e.g. a slice or integer # indexer. We dispatch to pyarrow. value = self._data[item] if isinstance(value, pa.ChunkedArray): return type(self)(value) else: scalar = value.as_py() if scalar is None: return self._dtype.na_value else: return scalar