SegArray optimization & bug fix

arkouda/join.py

@@ -117,7 +117,7 @@ def join_on_eq_with_dt(
     return resI, resJ
 
 
-def gen_ranges(starts, ends, stride=1):
+def gen_ranges(starts, ends, stride=1, return_lengths=False):
     """
     Generate a segmented array of variable-length, contiguous ranges between pairs of
     start- and end-points.
@@ -130,13 +130,17 @@ def gen_ranges(starts, ends, stride=1):
         The end value (exclusive) of each range
     stride: int
         Difference between successive elements of each range
+    return_lengths: bool, optional
+        Whether or not to return the lengths of each segment. Default False.
 
     Returns
     -------
     segments : pdarray, int64
         The starting index of each range in the resulting array
     ranges : pdarray, int64
         The actual ranges, flattened into a single array
+    lengths : pdarray, int64
+        The lengths of each segment. Only returned if return_lengths=True.
     """
     if starts.size != ends.size:
         raise ValueError("starts and ends must be same length")
@@ -158,7 +162,12 @@ def gen_ranges(starts, ends, stride=1):
         )
     )
     slices[segs[non_empty]] = diffs
-    return segs, cumsum(slices)
+
+    sums = cumsum(slices)
+    if return_lengths:
+        return segs, sums, lengths
+    else:
+        return segs, sums
 
 
 @typechecked

arkouda/segarray.py

@@ -96,7 +96,7 @@ def __init__(self, segments, values, lengths=None, grouping=None):
             self.lengths = lengths
 
         self._non_empty = self.lengths > 0
-        self._non_empty_count = self.non_empty.sum()
+        self._non_empty_count = self._non_empty.sum()
 
         # grouping object computation. (This will need to be moved to the server)
         # GroupBy computation left here because of lack of server obj. May need to move in Future
@@ -228,8 +228,8 @@ def __getitem__(self, i):
         elif (isinstance(i, pdarray) and i.dtype in [akint64, akuint64, akbool]) or isinstance(i, slice):
             starts = self.segments[i]
             ends = starts + self.lengths[i]
-            newsegs, inds = gen_ranges(starts, ends)
-            return SegArray(newsegs, self.values[inds])
+            newsegs, inds, lengths = gen_ranges(starts, ends, return_lengths=True)
+            return SegArray(newsegs, self.values[inds], lengths)
         else:
             raise TypeError(f"Invalid index type: {type(i)}")
 

benchmarks/dataframe.py

@@ -53,7 +53,6 @@ def time_ak_df_display(N_per_locale, trials, seed):
     timings = {op: [] for op in OPS}
     results = {}
     for i in range(trials):
-        timings = {op: [] for op in OPS}
         for op in timings.keys():
             fxn = getattr(df, op)
             start = time.time()