Merge pull request #179 from masonproffitt/fix_chunked_flatten_len

Fix len of ChunkedArray after chunksizes changes
scikit-hep · Aug 15, 2019 · 423ca48 · 423ca48
2 parents 635275c + ad68e41
commit 423ca48
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/awkward/array/chunked.py b/awkward/array/chunked.py
@@ -117,6 +117,7 @@ def knowchunksizes(self, until=None):
         until = min(until, len(self._chunks))
         for i in range(len(self._chunksizes), until):
             self._chunksizes.append(len(self._chunks[i]))
+            self._offsets = None
 
     def knowtype(self, at):
         if not 0 <= at < len(self._chunks):

diff --git a/tests/test_chunked.py b/tests/test_chunked.py
@@ -136,6 +136,11 @@ def test_chunked_get2d(self):
         assert a[[True, False, True, False, True, False, True, False, True, False], 0].tolist() == [0.0, 2.0, 4.0, 6.0, 8.0]
         assert a[[True, False, True, False, True, False, True, False, True, False], 1].tolist() == [0.0, 2.2, 4.4, 6.6, 8.8]
 
+    def test_chunked_flatten_len(self):
+        a = ChunkedArray([[[0, 1]], [[2, 3], [4, 5]], [[6, 7], [8, 9]]])
+        assert len(a) == 5
+        assert len(a.flatten()) == 10
+
     def test_appendable_append(self):
         a = AppendableArray(3, numpy.float64)
         assert a.tolist() == []