Repository: graphnet-team/graphnet
Issue: Issue within the parquet dataset class


The bug happens when trying to read or write a parquet file. It would seem that in parquet_dataset._query_table, it expects an event number which is then to be turned into a sequential index by referring to the index. However the index variable when I run it seems to already be a sequential index (not event number) the below change seems to be working for me, but I do not know about the possible knock on effect, I imagine that this function is called several places. 

def _query_table(
        self,
        table: str,
        columns: Union[List[str], str],
        index: int,
        selection: Optional[str] = None,
    ) -> List[Tuple[Any]]:
        # Check(s)
        assert (
            selection is None
        ), "Argument `selection` is currently not supported"

        sequential_index = index # replacing sequential_index =self._indices.index(index)   
         !!!!!!!!! ############## FIX WORKING FOR ME RIGHT HERE ################ !!!!!!!!!!!!!

        try:
            ak_array = self._parquet_hook[table][columns][sequential_index]
        except ValueError as e:
            if "does not exist (not in record)" in str(e):
                raise ColumnMissingException(str(e))
            else:
                raise e

        dictionary = ak_array.to_list()
        assert list(dictionary.keys()) == columns

        if all(map(np.isscalar, dictionary.values())):
            result = [tuple(dictionary.values())]

        else:
            # All arrays should have same length
            array_lengths = [
                len(values)
                for values in dictionary.values()
                if not np.isscalar(values)
            ]
            assert (
                len(set(array_lengths)) == 1
            ), f"Arrays in {dictionary} have differing lengths"
            nb_elements = array_lengths[0]

            # Broadcast scalars
            for key in dictionary:
                value = dictionary[key]
                if np.isscalar(value):
                    dictionary[key] = np.repeat(
                        value, repeats=nb_elements
                    ).tolist()

            result = list(map(tuple, list(zip(*dictionary.values()))))

        return result

**Expected behavior**
 sequential_index = self._indices[index]  was expected to take an index (possible by event_number) and turn it into a sequential index number

**Actual behavior**
recieves what I deem to be a sequential index number and returns error $index number$ not in index.




Changed files:
src/graphnet/data/parquet/parquet_dataset.py

diff --git a/src/graphnet/data/parquet/parquet_dataset.py b/src/graphnet/data/parquet/parquet_dataset.py
index 772d6c71..82caedaf 100644
--- a/src/graphnet/data/parquet/parquet_dataset.py
+++ b/src/graphnet/data/parquet/parquet_dataset.py
@@ -35,9 +35,13 @@ class ParquetDataset(Dataset):
         self._parquet_hook = ak.from_parquet(self._path, lazy=False)
 
     def _get_all_indices(self):
-        return ak.to_numpy(
-            self._parquet_hook[self._truth_table][self._index_column]
-        ).tolist()
+        return np.arange(
+            len(
+                ak.to_numpy(
+                    self._parquet_hook[self._truth_table][self._index_column]
+                ).tolist()
+            )
+        )
 
     def _query_table(
         self,
@@ -51,7 +55,7 @@ class ParquetDataset(Dataset):
             selection is None
         ), "Argument `selection` is currently not supported"
 
-        sequential_index = self._indices.index(index)
+        sequential_index = self._indices[index]
 
         try:
             ak_array = self._parquet_hook[table][columns][sequential_index]