Merge pull request #2288 from cta-observatory/fix_load_obsinfo

Only remove sort index from table if input did not contain it

ctapipe/io/tableloader.py

@@ -310,19 +310,23 @@ def _get_sort_index(self, start=None, stop=None):
  return table
 
  @staticmethod
- def _sort_to_original_order(table, include_tel_id=False):
+ def _sort_to_original_order(table, include_tel_id=False, keep_index=False):
  if len(table) == 0:
  return
  if include_tel_id:
  table.sort(("__index__", "tel_id"))
  else:
  table.sort("__index__")
- table.remove_column("__index__")
+
+ if not keep_index:
+ table.remove_column("__index__")
 
  @staticmethod
  def _add_index_if_needed(table):
  if "__index__" not in table.colnames:
  table["__index__"] = np.arange(len(table))
+ return True
+ return False
 
  def read_simulation_configuration(self):
  """
@@ -350,7 +354,7 @@ def _join_observation_info(self, table):
  obs_table["obs_id"] = obs_table["obs_id"].astype(table["obs_id"].dtype)
 
  # to be able to sort to original table order
- self._add_index_if_needed(table)
+ index_added = self._add_index_if_needed(table)
 
  joint = join_allow_empty(
  table,
@@ -360,8 +364,9 @@ def _join_observation_info(self, table):
  )
 
  # sort back to original order and remove index col
- self._sort_to_original_order(joint)
- del table["__index__"]
+ self._sort_to_original_order(joint, keep_index=not index_added)
+ if index_added:
+ table.remove_column("__index__")
  return joint
 
  def read_subarray_events(self, start=None, stop=None, keep_order=True):

ctapipe/io/tests/test_table_loader.py

@@ -415,13 +415,17 @@ def test_order_merged():
 
  path = get_dataset_path("gamma_diffuse_dl2_train_small.dl2.h5")
 
+ trigger = read_table(path, "/dl1/event/subarray/trigger")
  tel_trigger = read_table(path, "/dl1/event/telescope/trigger")
  with TableLoader(
  path,
- load_dl1_parameters=False,
+ load_dl1_parameters=True,
  load_dl2=True,
  load_observation_info=True,
  ) as loader:
+ events = loader.read_subarray_events()
+ check_equal_array_event_order(events, trigger)
+
  tables = loader.read_telescope_events_by_id()
 
  for tel_id, table in tables.items():

docs/changes/2288.bugfix.rst

@@ -0,0 +1,2 @@
+Fix ``TableLoader.read_subarray_events`` raising an exception when
+``load_observation_info=True``.