Skip to content

Commit

Permalink
fix: UnmaskedArray was missing its _nextcarry_outindex (#3097)
Browse files Browse the repository at this point in the history
* fix: UnmaskedArray was missing its _nextcarry_outindex

* add BitMaskedArray and ByteMaskedArray tests, too

---------

Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch>
  • Loading branch information
jpivarski and ianna committed May 3, 2024
1 parent 4773f24 commit ac26470
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/awkward/contents/unmaskedarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,12 @@ def _carry(self, carry: Index, allow_lazy: bool) -> Content:
self._content._carry(carry, allow_lazy), parameters=self._parameters
)

def _nextcarry_outindex(self) -> tuple[int, ak.index.Index64, ak.index.Index64]:
counting = self._backend.index_nplike.arange(self._content.length)
nextcarry = ak.index.Index64(counting, nplike=self._backend.index_nplike)
outindex = ak.index.Index64(counting, nplike=self._backend.index_nplike)
return 0, nextcarry, outindex

def _getitem_next_jagged(
self, slicestarts: Index, slicestops: Index, slicecontent: Content, tail
) -> Content:
Expand Down
52 changes: 52 additions & 0 deletions tests/test_3097_UnmaskedArray_missing_nextcarry_outindex.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

import numpy as np

import awkward as ak
from awkward.contents import (
BitMaskedArray,
ByteMaskedArray,
ListOffsetArray,
NumpyArray,
UnmaskedArray,
)
from awkward.index import Index8, Index64, IndexU8


def test_UnmaskedArray():
layout = ListOffsetArray(
Index64(np.array([0, 3, 3, 5])),
UnmaskedArray(NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5]))),
)
out = ak.drop_none(ak.Array(layout), axis=1)
assert out.tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]


def test_ByteMaskedArray():
layout = ListOffsetArray(
Index64(np.array([0, 3, 3, 5])),
ByteMaskedArray(
Index8(np.array([0, 1, 1, 0, 1])),
NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])),
True,
),
)
out = ak.drop_none(ak.Array(layout), axis=1)
assert out.tolist() == [[2.2, 3.3], [], [5.5]]


def test_BitMaskedArray():
layout = ListOffsetArray(
Index64(np.array([0, 3, 3, 5])),
BitMaskedArray(
IndexU8(np.array([15])),
NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])),
True,
5,
True,
),
)
out = ak.drop_none(ak.Array(layout), axis=1)
assert out.tolist() == [[1.1, 2.2, 3.3], [], [4.4]]

0 comments on commit ac26470

Please sign in to comment.