Skip to content

Commit

Permalink
Support boolean indexers and other properties in loc (#635)
Browse files Browse the repository at this point in the history
* Support boolean indexers and other properties in `loc`

* Resolves #610
* Adds a fastrack to lists of columns and boolean indexers

* Fix transposed getitem

* Address comments
  • Loading branch information
devin-petersohn authored and williamma12 committed May 26, 2019
1 parent 93fdadf commit 484d36c
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 3 deletions.
11 changes: 8 additions & 3 deletions modin/backends/pandas/query_compiler.py
Expand Up @@ -1925,11 +1925,13 @@ def back(self, n):
# End Head/Tail/Front/Back

# __getitem__ methods
def getitem_column_array(self, key):
def getitem_column_array(self, key, numeric=False):
"""Get column data for target labels.
Args:
key: Target labels by which to retrieve data.
numeric: A boolean representing whether or not the key passed in represents
the numeric index or the named index.
Returns:
A new QueryCompiler.
Expand All @@ -1941,7 +1943,10 @@ def getitem_column_array(self, key):
.transpose()
)
# Convert to list for type checking
numeric_indices = self.columns.get_indexer_for(key)
if not numeric:
numeric_indices = self.columns.get_indexer_for(key)
else:
numeric_indices = key
result = self.data.mask(col_indices=numeric_indices)
# We can't just set the columns to key here because there may be
# multiple instances of a key.
Expand All @@ -1962,7 +1967,7 @@ def getitem_row_array(self, key):
A new QueryCompiler.
"""
if self._is_transposed:
return self.transpose().getitem_column_array(key).transpose()
return self.transpose().getitem_column_array(key, numeric=True).transpose()
result = self.data.mask(row_indices=key)
# We can't just set the index to key here because there may be multiple
# instances of a key.
Expand Down
13 changes: 13 additions & 0 deletions modin/pandas/indexing.py
Expand Up @@ -227,6 +227,19 @@ class _LocIndexer(_LocationIndexerBase):
"""A indexer for ray_df.loc[] functionality"""

def __getitem__(self, key):
# When getting along a single axis,
if not isinstance(key, tuple):
# Try to fasttrack the code through already optimized path
try:
return self.df.__getitem__(key)
# This can happen if it is a list of rows
except KeyError:
pass
else:
if len(key) > self.df.ndim:
raise IndexingError("Too many indexers")
if key[0] == slice(None):
return self.df.__getitem__(key[1])
row_loc, col_loc, ndim, self.row_scaler, self.col_scaler = _parse_tuple(key)
self._handle_enlargement(row_loc, col_loc)
row_lookup, col_lookup = self._compute_lookup(row_loc, col_loc)
Expand Down

0 comments on commit 484d36c

Please sign in to comment.