Dense to sparse CSF conversion now in order of dimension size.

DataDog · Feb 5, 2020 · bd0d8c2 · bd0d8c2
1 parent eb51947
commit bd0d8c2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/cpp/src/arrow/sparse_tensor.cc b/cpp/src/arrow/sparse_tensor.cc
@@ -17,6 +17,7 @@
 
 #include "arrow/sparse_tensor.h"
 
+#include <arrow/util/sort.h>
 #include <algorithm>
 #include <functional>
 #include <limits>
@@ -454,18 +455,17 @@ class SparseTensorConverter<TYPE, SparseCSFIndex>
 
     std::vector<int64_t> counts(ndim);
     std::fill_n(counts.begin(), ndim, static_cast<int64_t>(0));
-
-    std::vector<int64_t> axis_order(ndim);
-    for (int64_t i = 0; i < ndim; ++i) axis_order[i] = i;
+    std::vector<int64_t> axis_order = internal::ArgSort(tensor_.shape());
 
     std::vector<TypedBufferBuilder<c_index_value_type>> indptr_buffer_builders(ndim - 1);
     std::vector<TypedBufferBuilder<c_index_value_type>> indices_buffer_builders(ndim);
 
     for (int64_t row = 0; row < nonzero_count; ++row) {
       bool tree_split = false;
       for (int64_t column = 0; column < ndim; ++column) {
-        bool change = coords->Value<IndexValueType>({row, column}) !=
-                      coords->Value<IndexValueType>({row - 1, column});
+        int64_t dimension = axis_order[column];
+        bool change = coords->Value<IndexValueType>({row, dimension}) !=
+                      coords->Value<IndexValueType>({row - 1, dimension});
 
         if (tree_split || change || row == 0) {
           if (row > 1) tree_split = true;
@@ -475,7 +475,7 @@ class SparseTensorConverter<TYPE, SparseCSFIndex>
                 static_cast<c_index_value_type>(counts[column + 1])));
           RETURN_NOT_OK(
               indices_buffer_builders[column].Append(static_cast<c_index_value_type>(
-                  coords->Value<IndexValueType>({row, column}))));
+                  coords->Value<IndexValueType>({row, dimension}))));
           ++counts[column];
         }
       }

diff --git a/format/SparseTensor.fbs b/format/SparseTensor.fbs
@@ -172,7 +172,7 @@ table SparseTensorIndexCSF {
   ///
   ///   axisOrder(X) = [0, 1, 2, 3].
   ///
-  axisOrder: [long];
+  axisOrder: [Int];
 }
 
 union SparseTensorIndex {