[NSE-417] Sort spill support framework (#369)

* adding unit test for sort spill Signed-off-by: Yuan Zhou <[email protected]> * fix format Signed-off-by: Yuan Zhou <[email protected]> * fix unit test Signed-off-by: Yuan Zhou <[email protected]> * support spill in sort Signed-off-by: Yuan Zhou <[email protected]> * fix read path Signed-off-by: Yuan Zhou <[email protected]> * support multiple keys Signed-off-by: Yuan Zhou <[email protected]> * update API for more kernels Signed-off-by: Yuan Zhou <[email protected]> * clean up Signed-off-by: Yuan Zhou <[email protected]> * support Spill API in sort Signed-off-by: Yuan Zhou <[email protected]> * reenable CoalesceBatch before Sort Signed-off-by: Yuan Zhou <[email protected]> * remove use cnt Signed-off-by: Yuan Zhou <[email protected]> * integrate with spark mem mgnt framework Signed-off-by: Yuan Zhou <[email protected]> * more spill to sort result iterator Signed-off-by: Yuan Zhou <[email protected]> * clean up Signed-off-by: Yuan Zhou <[email protected]> * release recordbatch from java size Signed-off-by: Yuan Zhou <[email protected]> * rebase Signed-off-by: Yuan Zhou <[email protected]> * fix unit test Signed-off-by: Yuan Zhou <[email protected]> * refine repeated spill Signed-off-by: Yuan Zhou <[email protected]> * fix format Signed-off-by: Yuan Zhou <[email protected]>
oap-project · Jul 23, 2021 · 5d05fd2 · 5d05fd2
1 parent 7168359
commit 5d05fd2
Show file tree

Hide file tree

Showing 23 changed files with 584 additions and 97 deletions.
diff --git a/native-sql-engine/core/src/main/java/com/intel/oap/vectorized/ExpressionEvaluator.java b/native-sql-engine/core/src/main/java/com/intel/oap/vectorized/ExpressionEvaluator.java
@@ -26,6 +26,7 @@
 import java.util.List;
 
 import org.apache.arrow.dataset.jni.NativeMemoryPool;
+import org.apache.arrow.dataset.jni.UnsafeRecordBatchSerializer;
 import org.apache.arrow.gandiva.evaluator.SelectionVectorInt16;
 import org.apache.arrow.gandiva.exceptions.GandivaException;
 import org.apache.arrow.gandiva.expression.ExpressionTree;
@@ -142,6 +143,25 @@ public ArrowRecordBatch[] evaluate(ArrowRecordBatch recordBatch) throws RuntimeE
     return evaluate(recordBatch, null);
   }
 
+  /**
+   * Evaluate input data using builded native function, and output as recordBatch.
+   */
+  public ArrowRecordBatch[] evaluate2(ArrowRecordBatch recordBatch) throws RuntimeException, IOException {
+    byte[] bytes = UnsafeRecordBatchSerializer.serializeUnsafe(recordBatch);
+    ArrowRecordBatchBuilder[] resRecordBatchBuilderList = jniWrapper.nativeEvaluate2(nativeHandler, bytes);
+    ArrowRecordBatch[] recordBatchList = new ArrowRecordBatch[resRecordBatchBuilderList.length];
+    for (int i = 0; i < resRecordBatchBuilderList.length; i++) {
+      if (resRecordBatchBuilderList[i] == null) {
+        recordBatchList[i] = null;
+        break;
+      }
+      ArrowRecordBatchBuilderImpl resRecordBatchBuilderImpl = new ArrowRecordBatchBuilderImpl(
+          resRecordBatchBuilderList[i]);
+      recordBatchList[i] = resRecordBatchBuilderImpl.build();
+    }
+    return recordBatchList;
+  }
+
   /**
    * Evaluate input data using builded native function, and output as recordBatch.
    */

diff --git a/...sql-engine/core/src/main/java/com/intel/oap/vectorized/ExpressionEvaluatorJniWrapper.java b/...sql-engine/core/src/main/java/com/intel/oap/vectorized/ExpressionEvaluatorJniWrapper.java
@@ -145,6 +145,8 @@ native long nativeBuildWithFinish(long memoryPool, byte[] schemaBuf, byte[] expr
         native ArrowRecordBatchBuilder[] nativeEvaluate(long nativeHandler, int numRows, long[] bufAddrs,
                         long[] bufSizes) throws RuntimeException;
 
+       native ArrowRecordBatchBuilder[] nativeEvaluate2(long nativeHandler, byte[] bytes) throws RuntimeException;
+
         /**
          * Get native kernel signature by the nativeHandler.
          *

diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/ColumnarPlugin.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/ColumnarPlugin.scala
@@ -112,8 +112,6 @@ case class ColumnarPreOverrides() extends Rule[SparkPlan] {
       val child = replaceWithColumnarPlan(plan.child)
       logDebug(s"Columnar Processing for ${plan.getClass} is currently supported.")
       child match {
-        case CoalesceBatchesExec(fwdChild: SparkPlan) =>
-          ColumnarSortExec(plan.sortOrder, plan.global, fwdChild, plan.testSpillFrequency)
         case _ =>
           ColumnarSortExec(plan.sortOrder, plan.global, child, plan.testSpillFrequency)
       }

diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarSorter.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarSorter.scala
@@ -92,8 +92,8 @@ class ColumnarSorter(
     elapse.set(NANOSECONDS.toMillis(total_elapse))
     sortTime.set(NANOSECONDS.toMillis(sort_elapse))
     shuffleTime.set(NANOSECONDS.toMillis(shuffle_elapse))
-    inputBatchHolder.foreach(cb => cb.close())
-    inputBatchHolder.clear
+    //inputBatchHolder.foreach(cb => cb.close())
+    //inputBatchHolder.clear
     if (sorter != null) {
       sorter.close()
     }
@@ -110,7 +110,7 @@ class ColumnarSorter(
     (0 until input.numCols).toList.foreach(i =>
       input.column(i).asInstanceOf[ArrowWritableColumnVector].retain())
     val beforeSort = System.nanoTime()
-    sorter.evaluate(input_batch)
+    sorter.evaluate2(input_batch)
     sort_elapse += System.nanoTime() - beforeSort
     total_elapse += System.nanoTime() - beforeSort
     ConverterUtils.releaseArrowRecordBatch(input_batch)
@@ -160,8 +160,8 @@ class ColumnarSorter(
           return true
         } else {
           has_next = false
-          inputBatchHolder.foreach(cb => cb.close())
-          inputBatchHolder.clear
+          //inputBatchHolder.foreach(cb => cb.close())
+          //inputBatchHolder.clear
           return false
         }
       }
@@ -373,7 +373,7 @@ object ColumnarSorter extends Logging {
     val (sort_expr, arrowSchema) = init(sortOrder, outputAttributes, sparkConf)
     val sorter = new ExpressionEvaluator()
     val signature = sorter
-      .build(arrowSchema, Lists.newArrayList(sort_expr), arrowSchema, true /*return at finish*/ )
+      .build(arrowSchema, Lists.newArrayList(sort_expr), arrowSchema, true /*return at finish*/)
     sorter.close
     signature
   }
@@ -391,7 +391,7 @@ object ColumnarSorter extends Logging {
     val (sort_expr, arrowSchema) = init(sortOrder, outputAttributes, sparkConf)
     val sorter = new ExpressionEvaluator(listJars.toList.asJava)
     sorter
-      .build(arrowSchema, Lists.newArrayList(sort_expr), arrowSchema, true /*return at finish*/ )
+      .build(arrowSchema, Lists.newArrayList(sort_expr), arrowSchema, true /*return at finish*/, true/*enable spill*/)
     new ColumnarSorter(
       sorter,
       outputAttributes,

diff --git a/native-sql-engine/cpp/compile.sh b/native-sql-engine/cpp/compile.sh
@@ -27,7 +27,7 @@ fi
 mkdir build
 cd build
 cmake .. -DTESTS=${TESTS} -DBUILD_ARROW=${BUILD_ARROW} -DSTATIC_ARROW=${STATIC_ARROW} -DBUILD_PROTOBUF=${BUILD_PROTOBUF} -DARROW_ROOT=${ARROW_ROOT} -DARROW_BFS_INSTALL_DIR=${ARROW_BFS_INSTALL_DIR}
-make
+make -j2
 
 set +eu
 
diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/code_generator.h b/native-sql-engine/cpp/src/codegen/arrow_compute/code_generator.h
@@ -114,7 +114,7 @@ class ArrowComputeCodeGenerator : public CodeGenerator {
     return arrow::Status::OK();
   }
 
-  arrow::Status evaluate(const std::shared_ptr<arrow::RecordBatch>& in,
+  arrow::Status evaluate(std::shared_ptr<arrow::RecordBatch>& in,
                          std::vector<std::shared_ptr<arrow::RecordBatch>>* out) {
     arrow::Status status = arrow::Status::OK();
     std::vector<ArrayList> batch_array;
@@ -256,14 +256,15 @@ class ArrowComputeCodeGenerator : public CodeGenerator {
     }
     int64_t current_spilled = 0L;
     for (auto visitor : visitor_list_) {
-      int64_t single_call_spilled;
+      int64_t single_call_spilled = 0;
       RETURN_NOT_OK(visitor->Spill(size - current_spilled, &single_call_spilled));
       current_spilled += single_call_spilled;
       if (current_spilled >= size) {
         *spilled_size = current_spilled;
         return arrow::Status::OK();
       }
     }
+
     *spilled_size = current_spilled;
     return arrow::Status::OK();
   }

diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor.cc b/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor.cc
@@ -517,8 +517,8 @@ arrow::Status ExprVisitor::Eval(const std::shared_ptr<arrow::Array>& selection_i
   return arrow::Status::OK();
 }
 
-arrow::Status ExprVisitor::Eval(const std::shared_ptr<arrow::RecordBatch>& in) {
-  in_record_batch_ = in;
+arrow::Status ExprVisitor::Eval(std::shared_ptr<arrow::RecordBatch>& in) {
+  in_record_batch_ = std::move(in);
   RETURN_NOT_OK(Eval());
   return arrow::Status::OK();
 }
@@ -757,19 +757,20 @@ arrow::Status ExprVisitor::GetResult(
 }
 
 arrow::Status ExprVisitor::Spill(int64_t size, int64_t* spilled_size) {
-  int64_t current_spilled = 0L;
+  int64_t current_spilled = 0;
   if (dependency_) {
     // fixme cycle invocation?
-    int64_t single_call_spilled;
+    int64_t single_call_spilled = 0;
     RETURN_NOT_OK(dependency_->Spill(size - current_spilled, &single_call_spilled));
     current_spilled += single_call_spilled;
+
     if (current_spilled >= size) {
       *spilled_size = current_spilled;
       return arrow::Status::OK();
     }
   }
   if (!finish_visitor_) {
-    int64_t single_call_spilled;
+    int64_t single_call_spilled = 0;
     RETURN_NOT_OK(impl_->Spill(size - current_spilled, &single_call_spilled));
     current_spilled += single_call_spilled;
   }

diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor.h b/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor.h
@@ -152,7 +152,7 @@ class ExprVisitor : public std::enable_shared_from_this<ExprVisitor> {
   arrow::Status Init();
   arrow::Status Eval(const std::shared_ptr<arrow::Array>& selection_in,
                      const std::shared_ptr<arrow::RecordBatch>& in);
-  arrow::Status Eval(const std::shared_ptr<arrow::RecordBatch>& in);
+  arrow::Status Eval(std::shared_ptr<arrow::RecordBatch>& in);
   arrow::Status Eval();
   std::string GetSignature() { return signature_; }
   arrow::Status SetMember(const std::shared_ptr<arrow::RecordBatch>& ms);
@@ -190,6 +190,7 @@ class ExprVisitor : public std::enable_shared_from_this<ExprVisitor> {
   std::shared_ptr<ExprVisitor> dependency_;
   std::shared_ptr<arrow::Array> in_selection_array_;
   std::shared_ptr<arrow::RecordBatch> in_record_batch_;
+  std::vector<std::shared_ptr<arrow::RecordBatch>> in_record_batch_holder_;
   std::vector<std::shared_ptr<arrow::Field>> ret_fields_;
 
   // For dual input kernels like probe

diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor_impl.h b/native-sql-engine/cpp/src/codegen/arrow_compute/expr_visitor_impl.h
@@ -496,8 +496,9 @@ class SortArraysToIndicesVisitorImpl : public ExprVisitorImpl {
       case ArrowComputeResultType::None: {
         std::vector<std::shared_ptr<arrow::Array>> col_list;
         for (auto col : p_->in_record_batch_->columns()) {
-          col_list.push_back(col);
+          col_list.push_back(std::move(col));
         }
+        p_->in_record_batch_holder_.push_back(std::move(p_->in_record_batch_));
         RETURN_NOT_OK(kernel_->Evaluate(col_list));
       } break;
       default:
@@ -508,6 +509,16 @@ class SortArraysToIndicesVisitorImpl : public ExprVisitorImpl {
     return arrow::Status::OK();
   }
 
+  arrow::Status Spill(int64_t size, int64_t* spilled_size) override {
+    std::cout << "target size: " << size << std::endl;
+    RETURN_NOT_OK(kernel_->Spill(size, spilled_size));
+
+    if (*spilled_size != 0) {
+      (p_->in_record_batch_holder_).clear();
+    }
+    return arrow::Status::OK();
+  }
+
   arrow::Status MakeResultIterator(std::shared_ptr<arrow::Schema> schema,
                                    std::shared_ptr<ResultIteratorBase>* out) override {
     switch (finish_return_type_) {