Using a list of indicies instead of single index in get(self, idx) in Dataset

[RasmusOrsoe]

Hi!

I have sqlite database that contains data that I'd like to query in the Dataset class. I have a list of sqlite indicies events (stored in self.selection' that I use to query the database with.

Conceptually, the get() function looks like

def get(self, idx):
      events = self.selection['event_no'][idx]
      with sqlite3.connect(self.database) as con:
          query = ¨select some variable from some table where events in %s' %events
          x_data = pd.read_sql(query, con)
      graphs = Data(x = torch.tensor(np.array(x_data), dtype = torch.float16))
      return graphs

However, this means that when the dataset is used in the ordinary way:

from torch_geometric.loader import DataLoader
dataset = MySQLiteDataset()
train_data = DataLoader(dataset, batch_size=1024, shuffle=True,num_workers=0)
batch = next(iter(train_data))

batch is constructed by querying the database 1024 times instead of a single time, which creates a quite big amount of computational overhead. Is there any way of forcing DataLoader to supply get(self, idx) with a batch of indicies and not a single index?

I have gone through both PyG and torch documentations but I just can't seem to figure this out. Any ideas are very welcome!

[rusty1s]

This sounds to me like you do not really need the Dataset class. Instead, you can do something like the following:

from torch_geometric.data import Data, Batch
from torch.utils.data import DataLoader

def collate_fn(indices: List[int]):
    # Query the database once
    # Creating a list of data objects from the query
    data_list = [Data(x=...), ...]
    return Batch.from_data_list(data_list)  # Create the batch

loader = DataLoader(range(num_graphs), batch_size=1024, collate_fn=collate_fn)

[RasmusOrsoe]

Thank you very much!