Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix slicing for UnmaskedArrays (which come from Arrow).
Browse files Browse the repository at this point in the history
jpivarski committed Jun 12, 2022
1 parent 1bdb85c commit e34292f
Showing 3 changed files with 35 additions and 4 deletions.
20 changes: 16 additions & 4 deletions src/awkward/_v2/_slicing.py
Original file line number Diff line number Diff line change
@@ -158,6 +158,7 @@ def prepare_tuple_RegularArray_toListOffsetArray64(item):

def prepare_tuple_nested(item):
if isinstance(item, ak._v2.contents.EmptyArray):
# policy: unknown -> int
return prepare_tuple_nested(item.toNumpyArray(np.int64))

elif isinstance(item, ak._v2.contents.NumpyArray) and issubclass(
@@ -217,7 +218,6 @@ def prepare_tuple_nested(item):
ak._v2.contents.UnmaskedArray,
),
):
# FIXME: might infinite-loop before simplify_optiontype is implemented
next = item.simplify_optiontype()
return prepare_tuple_nested(next)

@@ -258,8 +258,8 @@ def prepare_tuple_nested(item):
is_valid = item.mask_as_bool(valid_when=True)
positions_where_valid = item.nplike.index_nplike.nonzero(is_valid)[0]

nextcontent = prepare_tuple_nested(item.content)._carry(
ak._v2.index.Index64(positions_where_valid), False
nextcontent = prepare_tuple_nested(
item.content._carry(ak._v2.index.Index64(positions_where_valid), False)
)

nextindex = item.nplike.index_nplike.full(is_valid.shape[0], -1, np.int64)
@@ -275,7 +275,18 @@ def prepare_tuple_nested(item):
)

elif isinstance(item, ak._v2.contents.UnionArray):
return prepare_tuple_nested(item.simplify_uniontype())
attempt = item.simplify_uniontype()
if isinstance(attempt, ak._v2.contents.UnionArray):
raise ak._v2._util.error(
TypeError(
"irreducible unions (different types at the same level in an array) can't be used as slices"
)
)

return prepare_tuple_nested(attempt)

elif isinstance(item, ak._v2.contents.RecordArray):
raise ak._v2._util.error(TypeError("record arrays can't be used as slices"))

else:
raise ak._v2._util.error(
@@ -290,6 +301,7 @@ def prepare_tuple_nested(item):


def prepare_tuple_bool_to_int(item):
# actually convert leaf-node booleans to integers
if (
isinstance(item, ak._v2.contents.ListOffsetArray)
and isinstance(item.content, ak._v2.contents.NumpyArray)
9 changes: 9 additions & 0 deletions src/awkward/_v2/contents/emptyarray.py
Original file line number Diff line number Diff line change
@@ -103,6 +103,15 @@ def _carry(self, carry, allow_lazy):
else:
raise ak._v2._util.indexerror(self, carry.data, "array is empty")

def _getitem_next_jagged(self, slicestarts, slicestops, slicecontent, tail):
raise ak._v2._util.indexerror(
self,
ak._v2.contents.ListArray(
slicestarts, slicestops, slicecontent, None, None, self._nplike
),
"too many jagged slice dimensions for array",
)

def _getitem_next(self, head, tail, advanced):
if head == ():
return self
10 changes: 10 additions & 0 deletions src/awkward/_v2/contents/unmaskedarray.py
Original file line number Diff line number Diff line change
@@ -168,6 +168,16 @@ def _carry(self, carry, allow_lazy):
self._nplike,
)

def _getitem_next_jagged(self, slicestarts, slicestops, slicecontent, tail):
return UnmaskedArray(
self._content._getitem_next_jagged(
slicestarts, slicestops, slicecontent, tail
),
self._identifier,
self._parameters,
self._nplike,
)

def _getitem_next(self, head, tail, advanced):
if head == ():
return self

0 comments on commit e34292f

Please sign in to comment.