Skip to content

Commit

Permalink
Streamline recursively_apply for small slices of big arrays.
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpivarski
jpivarski committed May 2, 2022
1 parent fb6aa20 commit 25daa8f
Show file tree
Hide file tree
Showing 8 changed files with 94 additions and 21 deletions.
9 changes: 7 additions & 2 deletions src/awkward/_v2/contents/bitmaskedarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -543,12 +543,17 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if self._nplike.known_shape:
content = self._content[0 : self._length]
else:
content = self._content

if options["return_array"]:

def continuation():
return BitMaskedArray(
self._mask,
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand All @@ -566,7 +571,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand Down
9 changes: 7 additions & 2 deletions src/awkward/_v2/contents/bytemaskedarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -935,12 +935,17 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if self._nplike.known_shape:
content = self._content[0 : self._mask.length]
else:
content = self._content

if options["return_array"]:

def continuation():
return ByteMaskedArray(
self._mask,
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand All @@ -956,7 +961,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand Down
18 changes: 15 additions & 3 deletions src/awkward/_v2/contents/indexedarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1175,12 +1175,24 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if (
self._nplike.known_shape
and self._nplike.known_data
and self._index.length != 0
):
npindex = self._index.data
indexmin = npindex.min()
index = ak._v2.index.Index(npindex - indexmin, nplike=self._nplike)
content = self._content[indexmin : npindex.max() + 1]
else:
index, content = self._index, self._content

if options["return_array"]:

def continuation():
return IndexedArray(
self._index,
self._content._recursively_apply(
index,
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand All @@ -1195,7 +1207,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand Down
22 changes: 19 additions & 3 deletions src/awkward/_v2/contents/indexedoptionarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1563,12 +1563,28 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if (
self._nplike.known_shape
and self._nplike.known_data
and self._index.length != 0
):
npindex = self._index.data
npselect = npindex >= 0
if self._nplike.any(npselect):
indexmin = npindex[npselect].min()
index = ak._v2.index.Index(npindex - indexmin, nplike=self._nplike)
content = self._content[indexmin : npindex.max() + 1]
else:
index, content = self._index, self._content
else:
index, content = self._index, self._content

if options["return_array"]:

def continuation():
return IndexedOptionArray(
self._index,
self._content._recursively_apply(
index,
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand All @@ -1583,7 +1599,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth,
copy.copy(depth_context),
Expand Down
24 changes: 20 additions & 4 deletions src/awkward/_v2/contents/listarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1376,13 +1376,29 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if (
self._nplike.known_shape
and self._nplike.known_data
and self._starts.length != 0
):
startsmin = self._starts.data.min()
starts = ak._v2.index.Index(
self._starts.data - startsmin, nplike=self._nplike
)
stops = ak._v2.index.Index(
self._stops.data - startsmin, nplike=self._nplike
)
content = self._content[startsmin : self._stops.data.max()]
else:
starts, stops, content = self._starts, self._stops, self._content

if options["return_array"]:

def continuation():
return ListArray(
self._starts,
self._stops,
self._content._recursively_apply(
starts,
stops,
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand All @@ -1397,7 +1413,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand Down
15 changes: 12 additions & 3 deletions src/awkward/_v2/contents/listoffsetarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2020,12 +2020,21 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if self._nplike.known_shape and self._nplike.known_data:
offsetsmin = self._offsets.data[0]
offsets = ak._v2.index.Index(
self._offsets.data - offsetsmin, nplike=self._nplike
)
content = self._content[offsetsmin : self._offsets.data[-1]]
else:
offsets, content = self._offsets, self._content

if options["return_array"]:

def continuation():
return ListOffsetArray(
self._offsets,
self._content._recursively_apply(
offsets,
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand All @@ -2040,7 +2049,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand Down
9 changes: 7 additions & 2 deletions src/awkward/_v2/contents/recordarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -949,6 +949,11 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if self._nplike.known_shape:
contents = [x[: self._length] for x in self._contents]
else:
contents = self._contents

if options["return_array"]:

def continuation():
Expand All @@ -961,7 +966,7 @@ def continuation():
lateral_context,
options,
)
for content in self._contents
for content in contents
],
self._fields,
self._length,
Expand All @@ -973,7 +978,7 @@ def continuation():
else:

def continuation():
for content in self._contents:
for content in contents:
content._recursively_apply(
action,
depth,
Expand Down
9 changes: 7 additions & 2 deletions src/awkward/_v2/contents/regulararray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1153,11 +1153,16 @@ def _completely_flatten(self, nplike, options):
def _recursively_apply(
self, action, depth, depth_context, lateral_context, options
):
if self._nplike.known_shape:
content = self._content[: self._length * self._size]
else:
content = self._content

if options["return_array"]:

def continuation():
return RegularArray(
self._content._recursively_apply(
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand All @@ -1174,7 +1179,7 @@ def continuation():
else:

def continuation():
self._content._recursively_apply(
content._recursively_apply(
action,
depth + 1,
copy.copy(depth_context),
Expand Down

0 comments on commit 25daa8f

Please sign in to comment.