Fix ArrowDecoder.decode to return instead of yield (#2976)

awslabs · Aug 22, 2023 · c11f91f · c11f91f
1 parent db14c0f
commit c11f91f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/gluonts/dataset/arrow/dec.py b/src/gluonts/dataset/arrow/dec.py
@@ -32,7 +32,7 @@ def from_schema(cls, schema):
         )
 
     def decode(self, batch, row_number: int):
-        yield from self.decode_batch(batch.slice(row_number, row_number + 1))
+        return next(self.decode_batch(batch.slice(row_number, row_number + 1)))
 
     def decode_batch(self, batch):
         for row in batch.to_pandas().to_dict("records"):

diff --git a/test/dataset/test_arrow.py b/test/dataset/test_arrow.py
@@ -82,6 +82,8 @@ def test_arrow(writer, flatten_arrays):
 
         assert_equal(data, dataset)
 
+        assert_equal(dataset[4], data[4])
+
         assert len(dataset[:5]) == len(data[:5])
         assert_equal(dataset[:5], data[:5])