Closes #3870: bug in reshape for bigint type

Bears-R-Us · Nov 21, 2024 · d05990e · d05990e
1 parent b456d3e
commit d05990e
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 8 deletions.
diff --git a/src/AryUtil.chpl b/src/AryUtil.chpl
@@ -15,6 +15,7 @@ module AryUtil
     use List;
     use CommAggregation;
     use CommPrimitives;
+    use BigInteger;
 
 
     param bitsPerDigit = RSLSD_bitsPerDigit;
@@ -905,7 +906,8 @@ module AryUtil
     /*
       unflatten a 1D array into a multi-dimensional array of the given shape
     */
-    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws {
+    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws
+      where t!=bigint {
       var unflat = makeDistArray((...shape), t);
 
       if N == 1 {
@@ -952,7 +954,6 @@ module AryUtil
             // flat region is spread across multiple locales, do a get for each source locale
             for locInID in locInStart..locInStop {
               const flatSubSlice = flatSlice[flatLocRanges[locInID]];
-
               get(
                 c_ptrTo(unflat[dufc.orderToIndex(flatSubSlice.low)]),
                 getAddr(a[flatSubSlice.low]),
@@ -967,11 +968,30 @@ module AryUtil
       return unflat;
     }
 
+    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws
+      where t==bigint {
+      var unflat = makeDistArray((...shape), t);
+
+      if N == 1 {
+        unflat = a;
+        return unflat;
+      }
+
+      coforall loc in Locales with (ref unflat) do on loc {
+        forall idx in a.localSubdomain() with (var agg = newDstAggregator(t)) {
+          agg.copy(unflat[unflat.domain.orderToIndex(idx)], a[idx]);
+        }
+      }
+
+      return unflat;
+    }
+
     /*
       flatten a multi-dimensional array into a 1D array
     */
-    @arkouda.registerCommand
-    proc flatten(const ref a: [?d] ?t): [] t throws {
+    @arkouda.registerCommand(ignoreWhereClause=true)
+    proc flatten(const ref a: [?d] ?t): [] t throws
+      where t!=bigint {
       if a.rank == 1 then return a;
 
       var flat = makeDistArray(d.size, t);
@@ -1030,6 +1050,22 @@ module AryUtil
       return flat;
     }
 
+
+    proc flatten(const ref a: [?d] ?t): [] t throws
+      where t==bigint {
+      if a.rank == 1 then return a;
+
+      var flat = makeDistArray(d.size, t);
+
+      coforall loc in Locales with (ref flat) do on loc {
+        forall idx in flat.localSubdomain() with (var agg = newSrcAggregator(t)) {
+          agg.copy(flat[idx], a[a.domain.orderToIndex(idx)]);
+        }
+      }
+
+      return flat;
+    }
+
     // helper for computing an array element's index from its order
     record orderer {
       param rank: int;
@@ -1043,8 +1079,8 @@ module AryUtil
 
       // index -> order for the input array's indices
       // e.g., order = k + (nz * j) + (nz * ny * i)
-      inline proc indexToOrder(idx: rank*?t): t
-      where (t==int) || (t==uint(64)) {
+      inline proc indexToOrder(idx: rank*int): int
+        where (t==int) || (t==uint(64)) {
         var order : t = 0;
         for param i in 0..<rank do order += idx[i] * accumRankSizes[rank - i - 1];
         return order;

diff --git a/tests/pdarrayclass_test.py b/tests/pdarrayclass_test.py
@@ -23,12 +23,20 @@
 class TestPdarrayClass:
 
     @pytest.mark.skip_if_max_rank_less_than(2)
-    def test_reshape(self):
-        a = ak.arange(4)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_reshape(self, dtype):
+        a = ak.arange(4, dtype=dtype)
         r = a.reshape((2, 2))
         assert r.shape == (2, 2)
         assert isinstance(r, ak.pdarray)
 
+    @pytest.mark.skip_if_max_rank_less_than(3)
+    def test_reshape_bug_reproducer(self):
+        dtype = "bigint"
+        size = 10
+        x = ak.arange(size, dtype=dtype).reshape((1, size, 1))
+        ak_assert_equal(x.flatten(), ak.arange(size, dtype=dtype))
+
     def test_shape(self):
         a = ak.arange(4)
         np_a = np.arange(4)