Skip to content

Commit

Permalink
BUG: fix a crash when retrieving rows from a multiple-index table
Browse files Browse the repository at this point in the history
  • Loading branch information
neutrinoceros committed Jan 8, 2024
1 parent 5340e31 commit 3147175
Showing 1 changed file with 16 additions and 5 deletions.
21 changes: 16 additions & 5 deletions astropy/table/index.py
Expand Up @@ -30,14 +30,19 @@
c[1:2] -> reference
array.view(Column) -> no indices
"""
from __future__ import annotations

from copy import deepcopy
from typing import TYPE_CHECKING

import numpy as np

from .bst import MaxValue, MinValue
from .sorted_array import SortedArray

if TYPE_CHECKING:
from collections.abc import Iterable


class QueryError(ValueError):
"""
Expand Down Expand Up @@ -829,18 +834,24 @@ def __init__(self, table):
if len(self.indices) == 0:
raise ValueError("Cannot create TableLoc object with no indices")

def _get_rows(self, item):
def _get_rows(
self,
item: int | Iterable[int] | slice | tuple[str, slice],
) -> list[np.int64]:
"""
Retrieve Table rows indexes by value slice.
"""
if isinstance(item, tuple):
if (

Check warning on line 844 in astropy/table/index.py

View check run for this annotation

Codecov / codecov/patch

astropy/table/index.py#L844

Added line #L844 was not covered by tests
isinstance(item, tuple)
and len(item) == 2
and isinstance(item[0], str)
and isinstance(item[1], slice)
):
key, item = item
else:
key = self.table.primary_key

index = self.indices[key]
if len(index.columns) > 1:
raise ValueError("Cannot use .loc on multi-column indices")

if isinstance(item, slice):
# None signifies no upper/lower bound
Expand All @@ -853,7 +864,7 @@ def _get_rows(self, item):
# item should be a list or ndarray of values
rows = []
for key in item:
p = index.find((key,))
p = index.find(key if isinstance(key, tuple) else (key,))

Check warning on line 867 in astropy/table/index.py

View check run for this annotation

Codecov / codecov/patch

astropy/table/index.py#L867

Added line #L867 was not covered by tests
if len(p) == 0:
raise KeyError(f"No matches found for key {key}")
else:
Expand Down

0 comments on commit 3147175

Please sign in to comment.