Skip to content

Commit

Permalink
Fix negative integer indexing for Table
Browse files Browse the repository at this point in the history
  • Loading branch information
moeyensj committed Aug 14, 2023
1 parent 59e99ce commit 4394bfc
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 0 deletions.
2 changes: 2 additions & 0 deletions quivr/tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,8 @@ def __getitem__(self, idx: Union[int, slice]) -> Self:
:param idx: The row index or slice to return.
"""
if isinstance(idx, int):
if idx < 0:
idx += len(self)
table = self.table[idx : idx + 1]
else:
table = self.table[idx]
Expand Down
20 changes: 20 additions & 0 deletions test/test_tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,26 @@ class B(qv.Table):
assert b_table != b_table_attr_diff


def test_table__getitem__():
# Test indexing with positive and negative integers
table = Pair.from_kwargs(
x=[1, 2, 3],
y=[4, 5, 6],
)
assert table[0] == Pair.from_kwargs(x=1, y=4)
assert table[1] == Pair.from_kwargs(x=2, y=5)
assert table[2] == Pair.from_kwargs(x=3, y=6)
assert table[-1] == Pair.from_kwargs(x=3, y=6)
assert table[-2] == Pair.from_kwargs(x=2, y=5)
assert table[-3] == Pair.from_kwargs(x=1, y=4)

# Test indexing with slices
assert table[0:2] == Pair.from_kwargs(x=[1, 2], y=[4, 5])
assert table[1:3] == Pair.from_kwargs(x=[2, 3], y=[5, 6])
assert table[-2:] == Pair.from_kwargs(x=[2, 3], y=[5, 6])
assert table[:-1] == Pair.from_kwargs(x=[1, 2], y=[4, 5])


def test_table_to_structarray():
xs = pa.array([1, 2, 3], pa.int64())
ys = pa.array([4, 5, 6], pa.int64())
Expand Down

0 comments on commit 4394bfc

Please sign in to comment.