Skip to content

Commit

Permalink
fix: don't project records during broadcasting; push index down (#2524)
Browse files Browse the repository at this point in the history
  • Loading branch information
agoose77 committed Jun 14, 2023
1 parent f81d2bf commit 23dc2ae
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/awkward/_broadcasting.py
Original file line number Diff line number Diff line change
Expand Up @@ -934,7 +934,13 @@ def broadcast_any_union():
)

def broadcast_any_indexed():
nextinputs = [x.project() if isinstance(x, IndexedArray) else x for x in inputs]
# The `apply` function may exit at the level of a `RecordArray`. We can avoid projection
# of the record array in such cases, in favour of a deferred carry. This can be done by
# "pushing" the `IndexedArray` _into_ the record (i.e., wrapping each `content`).
nextinputs = [
x._push_inside_record_or_project() if isinstance(x, IndexedArray) else x
for x in inputs
]
return apply_step(
backend,
nextinputs,
Expand Down
15 changes: 15 additions & 0 deletions src/awkward/contents/indexedarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1136,3 +1136,18 @@ def _is_equal_to(self, other, index_dtype, numpyarray):
return self.index.is_equal_to(
other.index, index_dtype, numpyarray
) and self.content.is_equal_to(other.content, index_dtype, numpyarray)

def _push_inside_record_or_project(self) -> Self | ak.contents.RecordArray:
if self.content.is_record:
return ak.contents.RecordArray(
contents=[
ak.contents.IndexedArray.simplified(self._index, c)
for c in self.content.contents
],
fields=self.content._fields,
length=self.length,
backend=self._backend,
parameters=parameters_union(self.content._parameters, self._parameters),
)
else:
return self.project()

0 comments on commit 23dc2ae

Please sign in to comment.