fix: dataset.concat reworked to allow operations like subset and merge afterwards

Author: samedii

.flake8

@@ -0,0 +1,4 @@
+[flake8]
+max-line-length = 80
+select = C,E,F,W,B,B950
+extend-ignore = E203, E501

datastream/dataset.py

@@ -5,6 +5,8 @@
 )
 from pathlib import Path
 from functools import lru_cache
+import string
+import random
 import textwrap
 import inspect
 import numpy as np
@@ -377,13 +379,23 @@ def get_item(dataframe, index):
                 get_item=get_item,
             )
         else:
+            dataset_column = (
+                '__concat__'
+                + ''.join([random.choice(string.ascii_lowercase) for _ in range(8)])
+            )
+
+            new_dataframe = pd.concat([dataset.dataframe for dataset in datasets])
+            new_dataframe[dataset_column] = [
+                from_concat_mapping(index)[0]
+                for index in range(len(new_dataframe))
+            ]
 
             def get_item(dataframe, index):
-                dataset_index, _ = from_concat_mapping(index)
+                dataset_index = int(dataframe.iloc[index][dataset_column])
                 return datasets[dataset_index].get_item(dataframe, index)
 
             return Dataset(
-                dataframe=pd.concat([dataset.dataframe for dataset in datasets]),
+                dataframe=new_dataframe,
                 length=sum(map(len, datasets)),
                 get_item=get_item,
             )

datastream/datastream.py

@@ -527,7 +527,7 @@ def test_sequential_sampler():
     assert next(it) == ['a', 'b', 'c', 'a', 'b', 'c']
 
 
-def test_merge_concat():
+def test_concat_merge():
     dataset = Dataset.concat([
         Dataset.from_subscriptable([1, 2]),
         Dataset.from_subscriptable([1, 3, 5]),
@@ -536,8 +536,51 @@ def test_merge_concat():
     datastream = Datastream.merge([
         Datastream(dataset),
         Datastream(dataset.subset(
-            lambda df: df["index"] <= 3
+            lambda df: [index < 3 for index in range(len(df))]
         )),
     ])
 
-    list(datastream)
+    assert len(dataset.subset(
+        lambda df: [index < 3 for index in range(len(df))]
+    )) == 3
+
+    assert len(list(datastream)) == 6
+
+
+def test_combine_concat_merge():
+    dataset = Dataset.concat([
+        Dataset.zip([
+            Dataset.from_subscriptable([1]),
+            Dataset.from_subscriptable([2]),
+        ]),
+        Dataset.combine([
+            Dataset.from_subscriptable([3, 3]),
+            Dataset.from_subscriptable([4, 4, 4]),
+        ]),
+    ])
+
+    datastream = Datastream.merge([
+        Datastream(dataset),
+        Datastream(Dataset.zip([
+            Dataset.from_subscriptable([5]),
+            Dataset.from_subscriptable([6]),
+        ])),
+    ])
+
+    assert len(list(datastream)) == 2
+
+
+def test_last_batch():
+    from datastream.samplers import SequentialSampler
+
+    datastream = Datastream(
+        Dataset.from_subscriptable(list('abc'))
+    )
+    assert list(map(len, datastream.data_loader(batch_size=4))) == [3]
+    assert list(map(len, datastream.data_loader(batch_size=4, n_batches_per_epoch=2))) == [4, 4]
+
+    datastream = Datastream(
+        Dataset.from_subscriptable(list('abc')),
+        SequentialSampler(3),
+    )
+    assert list(map(len, datastream.data_loader(batch_size=2))) == [2, 1]