in torcharrow/icolumn.py [0:0]
def __getitem__(self, arg):
"""
If *arg* is a
`n`, a number, return the row with index n
`[n1,..,nm]` return a new column with the rows[n1],..,rows[nm]
`[n1:n2:n3]`, return a new column slice with rows[n1:n2:n3]
`s`, a string, return the column named s
`[s1,..,sm]` return dataframe having column[s1],..,column[sm]
`[s1:s2]` return dataframe having columns[s1:s2]
`[b1,..,bn]`, where bi are booleans, return all rows that are true
`Column([b1..bn])` return all rows that are true
"""
if isinstance(arg, int):
return self._get(arg)
elif isinstance(arg, str):
return self._get_column(arg)
elif isinstance(arg, slice):
args = []
for i in [arg.start, arg.stop, arg.step]:
if isinstance(i, np.integer):
args.append(int(i))
else:
args.append(i)
if all(a is None or isinstance(a, int) for a in args):
return self._slice(*args)
else:
raise TypeError(
f"slice arguments {[type(a) for a in args]} should all be int"
)
elif isinstance(arg, list):
if len(arg) == 0:
return ta.DataFrame(device=self.device)
if all(isinstance(a, bool) for a in arg):
return self.filter(arg)
if all(isinstance(a, int) for a in arg):
return self._gets(arg)
if all(isinstance(a, str) for a in arg):
return self._get_columns(arg)
else:
raise TypeError("index should be list of int or list of str")
elif isinstance(arg, IColumn) and dt.is_boolean(arg.dtype):
return self.filter(arg)
else:
raise self._not_supported("__getitem__")