Merge pull request #686 from RasmusOrsoe/parquet_hangtime_bugfix

Bugfix to #685 and #683

src/graphnet/data/datamodule.py

@@ -1,7 +1,6 @@
 """Base `Dataloader` class(es) used in `graphnet`."""
 from typing import Dict, Any, Optional, List, Tuple, Union
 import pytorch_lightning as pl
-from torch.utils.data import DataLoader
 from copy import deepcopy
 from sklearn.model_selection import train_test_split
 import pandas as pd
@@ -13,6 +12,7 @@
     ParquetDataset,
 )
 from graphnet.utilities.logging import Logger
+from graphnet.data.dataloader import DataLoader
 
 
 class GraphNeTDataModule(pl.LightningDataModule, Logger):

src/graphnet/data/dataset/parquet/parquet_dataset.py

@@ -284,8 +284,10 @@ def _query_table(
             )
             data = df.select(columns)
             if isinstance(data[columns[0]][0], Series):
-                data = data.explode(columns)
-            array = data.to_numpy()
+                x = [data[col][0].to_numpy().reshape(-1, 1) for col in columns]
+                array = np.concatenate(x, axis=1)
+            else:
+                array = data.to_numpy()
         else:
             array = np.array()
         return array