diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ColumnarShuffledHashJoinExec.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ColumnarShuffledHashJoinExec.scala index 35a3de30f..59b9acfd1 100644 --- a/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ColumnarShuffledHashJoinExec.scala +++ b/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ColumnarShuffledHashJoinExec.scala @@ -356,31 +356,47 @@ case class ColumnarShuffledHashJoinExec( // now we can return this wholestagecodegen iter val res = new Iterator[ColumnarBatch] { - override def hasNext: Boolean = iter.hasNext + override def hasNext: Boolean = { + nativeIterator.hasNext || iter.hasNext + } override def next(): ColumnarBatch = { - val cb = iter.next() - val beforeEval = System.nanoTime() - val input_rb = - ConverterUtils.createArrowRecordBatch(cb) - if (input_rb.getLength == 0) { - ConverterUtils.releaseArrowRecordBatch(input_rb) + if (nativeIterator.hasNext) { + val beforeEval = System.nanoTime() + val out_batch = nativeIterator.next + val output = ConverterUtils.fromArrowRecordBatch(probe_out_schema, out_batch) + ConverterUtils.releaseArrowRecordBatch(out_batch) eval_elapse += System.nanoTime() - beforeEval - return createEmptyBatch() - } - val output_rb = nativeIterator.process(probe_input_schema, input_rb) - if (output_rb == null) { - ConverterUtils.releaseArrowRecordBatch(input_rb) + val outputNumRows = out_batch.getLength + numOutputRows += outputNumRows + val out_cb = new ColumnarBatch(output.map(v => v.asInstanceOf[ColumnVector]).toArray, outputNumRows) + return out_cb + } else { + val cb = iter.next() + val beforeEval = System.nanoTime() + val input_rb = + ConverterUtils.createArrowRecordBatch(cb) + if (input_rb.getLength == 0) { + ConverterUtils.releaseArrowRecordBatch(input_rb) + eval_elapse += System.nanoTime() - beforeEval + return createEmptyBatch() + } + val output_rb = nativeIterator.process(probe_input_schema, input_rb) + if (output_rb == null || output_rb.getLength == 0) { + ConverterUtils.releaseArrowRecordBatch(input_rb) + eval_elapse += System.nanoTime() - beforeEval + return createEmptyBatch() + } + + val out_batch = nativeIterator.next + val output = ConverterUtils.fromArrowRecordBatch(probe_out_schema, out_batch) + ConverterUtils.releaseArrowRecordBatch(out_batch) + val outputNumRows = out_batch.getLength eval_elapse += System.nanoTime() - beforeEval - return createEmptyBatch() + numOutputRows += outputNumRows + val out_cb = new ColumnarBatch(output.map(v => v.asInstanceOf[ColumnVector]).toArray, outputNumRows) + out_cb } - val outputNumRows = output_rb.getLength - ConverterUtils.releaseArrowRecordBatch(input_rb) - val output = ConverterUtils.fromArrowRecordBatch(probe_out_schema, output_rb) - ConverterUtils.releaseArrowRecordBatch(output_rb) - eval_elapse += System.nanoTime() - beforeEval - numOutputRows += outputNumRows - new ColumnarBatch(output.map(v => v.asInstanceOf[ColumnVector]).toArray, outputNumRows) } private def createEmptyBatch() = { diff --git a/native-sql-engine/cpp/src/codegen/arrow_compute/ext/conditioned_probe_kernel.cc b/native-sql-engine/cpp/src/codegen/arrow_compute/ext/conditioned_probe_kernel.cc index e30451578..85fba9496 100644 --- a/native-sql-engine/cpp/src/codegen/arrow_compute/ext/conditioned_probe_kernel.cc +++ b/native-sql-engine/cpp/src/codegen/arrow_compute/ext/conditioned_probe_kernel.cc @@ -16,6 +16,7 @@ */ #include +#include #include #include #include @@ -349,6 +350,7 @@ class ConditionedProbeKernel::Impl { left_field_list_(left_field_list), right_field_list_(right_field_list) { result_schema_ = arrow::schema(result_schema); + batch_size_ = GetBatchSize(); auto configuration = gandiva::ConfigurationBuilder().DefaultConfiguration(); for (auto expr : right_key_project_list[1]) { @@ -514,7 +516,40 @@ class ConditionedProbeKernel::Impl { } RETURN_NOT_OK(appender->Reset()); } - *out = arrow::RecordBatch::Make(result_schema_, out_length, out_arr_list); + *out = record_batch_holder_ = + arrow::RecordBatch::Make(result_schema_, out_length, out_arr_list); + // Initialize for iterator + record_batch_holder_length_ = out_length; + record_batch_holder_offset_ = 0; + return arrow::Status::OK(); + } + + bool HasNext() { + if (record_batch_holder_offset_ >= record_batch_holder_length_) { + return false; + } + return true; + } + + arrow::Status Next(std::shared_ptr* out) { + if (record_batch_holder_length_ <= batch_size_) { + *out = record_batch_holder_; + record_batch_holder_offset_ = record_batch_holder_length_; + return arrow::Status::OK(); + } + auto length = + (record_batch_holder_length_ - record_batch_holder_offset_) > batch_size_ + ? batch_size_ + : (record_batch_holder_length_ - record_batch_holder_offset_); + std::vector> out_arrs; + for (auto arr : record_batch_holder_->columns()) { + std::shared_ptr copied; + auto sliced = arr->Slice(record_batch_holder_offset_, length); + RETURN_NOT_OK(arrow::Concatenate({sliced}, ctx_->memory_pool(), &copied)); + out_arrs.push_back(copied); + } + *out = arrow::RecordBatch::Make(result_schema_, length, out_arrs); + record_batch_holder_offset_ += length; return arrow::Status::OK(); } @@ -1339,6 +1374,11 @@ class ConditionedProbeKernel::Impl { gandiva::FieldVector right_field_list_; gandiva::FieldVector right_projected_field_list_; std::shared_ptr probe_func_; + + std::shared_ptr record_batch_holder_; + uint64_t batch_size_; + uint64_t record_batch_holder_length_ = 0; + uint64_t record_batch_holder_offset_ = 0; }; arrow::Status GetInnerJoin(bool cond_check, std::string index_name, diff --git a/native-sql-engine/cpp/src/tests/arrow_compute_test_join_wocg.cc b/native-sql-engine/cpp/src/tests/arrow_compute_test_join_wocg.cc index bb6a603a9..4b67cb7fb 100644 --- a/native-sql-engine/cpp/src/tests/arrow_compute_test_join_wocg.cc +++ b/native-sql-engine/cpp/src/tests/arrow_compute_test_join_wocg.cc @@ -1093,6 +1093,174 @@ TEST(TestArrowComputeWSCG, JoinWOCGTestTwoStringInnerJoinType2) { } } +TEST(TestArrowComputeWSCG, JoinTestUsingInnerJoinwith2KeysforHashProbeIterator) { + setenv("NATIVESQL_BATCH_SIZE", "2", 1); + ////////////////////// prepare expr_vector /////////////////////// + auto table0_f0 = field("table0_f0", utf8()); + auto table0_f1 = field("table0_f1", utf8()); + auto table0_f2 = field("table0_f2", uint32()); + auto table1_f0 = field("table1_f0", utf8()); + auto table1_f1 = field("table1_f1", utf8()); + + /////////////////////////////////////////// + auto n_left = TreeExprBuilder::MakeFunction( + "codegen_left_schema", + {TreeExprBuilder::MakeField(table0_f0), TreeExprBuilder::MakeField(table0_f1), + TreeExprBuilder::MakeField(table0_f2)}, + uint32()); + auto n_right = TreeExprBuilder::MakeFunction( + "codegen_right_schema", + {TreeExprBuilder::MakeField(table1_f0), TreeExprBuilder::MakeField(table1_f1)}, + uint32()); + auto f_res = field("res", uint32()); + + auto n_left_key = TreeExprBuilder::MakeFunction( + "codegen_left_key_schema", + {TreeExprBuilder::MakeField(table0_f0), TreeExprBuilder::MakeField(table0_f1)}, + uint32()); + auto n_right_key = TreeExprBuilder::MakeFunction( + "codegen_right_key_schema", + {TreeExprBuilder::MakeField(table1_f0), TreeExprBuilder::MakeField(table1_f1)}, + uint32()); + auto n_result = TreeExprBuilder::MakeFunction( + "result", + {TreeExprBuilder::MakeField(table0_f0), TreeExprBuilder::MakeField(table0_f1), + TreeExprBuilder::MakeField(table0_f2), TreeExprBuilder::MakeField(table1_f0), + TreeExprBuilder::MakeField(table1_f1)}, + uint32()); + auto n_hash_config = TreeExprBuilder::MakeFunction( + "build_keys_config_node", {TreeExprBuilder::MakeLiteral((int)1)}, uint32()); + auto n_probeArrays = TreeExprBuilder::MakeFunction( + "conditionedProbeArraysInner", + {n_left, n_right, n_left_key, n_right_key, n_result, n_hash_config}, uint32()); + auto n_standalone = + TreeExprBuilder::MakeFunction("standalone", {n_probeArrays}, uint32()); + auto probeArrays_expr = TreeExprBuilder::MakeExpression(n_standalone, f_res); + + auto schema_table_0 = arrow::schema({table0_f0, table0_f1, table0_f2}); + auto schema_table_1 = arrow::schema({table1_f0, table1_f1}); + auto schema_table = + arrow::schema({table0_f0, table0_f1, table0_f2, table1_f0, table1_f1}); + + auto n_hash_kernel = TreeExprBuilder::MakeFunction( + "HashRelation", {n_left_key, n_hash_config}, uint32()); + auto n_hash = TreeExprBuilder::MakeFunction("standalone", {n_hash_kernel}, uint32()); + auto hashRelation_expr = TreeExprBuilder::MakeExpression(n_hash, f_res); + std::shared_ptr expr_build; + arrow::compute::ExecContext ctx; + ASSERT_NOT_OK(CreateCodeGenerator(ctx.memory_pool(), schema_table_0, + {hashRelation_expr}, {}, &expr_build, true)); + std::shared_ptr expr_probe; + ASSERT_NOT_OK(CreateCodeGenerator( + ctx.memory_pool(), schema_table_1, {probeArrays_expr}, + {table0_f0, table0_f1, table0_f2, table1_f0, table1_f1}, &expr_probe, true)); + ///////////////////// Calculation ////////////////// + std::shared_ptr input_batch; + + std::vector> dummy_result_batches; + + std::vector> table_0; + std::vector> table_1; + + std::vector input_data_string = { + R"(["l", "c", "a", "b"])", R"(["L", "C", "A", "B"])", "[10, 3, 1, 2]"}; + MakeInputBatch(input_data_string, schema_table_0, &input_batch); + table_0.push_back(input_batch); + + input_data_string = {R"(["f", "n", "e", "j"])", R"(["F", "N", "E", "J"])", + "[6, 12, 5, 8]"}; + MakeInputBatch(input_data_string, schema_table_0, &input_batch); + table_0.push_back(input_batch); + + std::vector input_data_2_string = {R"(["a", "b", "c", "d", "e", "f"])", + R"(["A", "B", "C", "D", "F", "F"])"}; + MakeInputBatch(input_data_2_string, schema_table_1, &input_batch); + table_1.push_back(input_batch); + + input_data_2_string = {R"(["i", "j", "k", "l", "m", "n"])", + R"(["I", "J", "K", "L", "M", "N"])"}; + MakeInputBatch(input_data_2_string, schema_table_1, &input_batch); + table_1.push_back(input_batch); + + //////////////////////// data prepared ///////////////////////// + + std::vector> expected_table; + std::shared_ptr expected_result; + std::vector expected_result_string = { + R"(["a", "b", "c", "f"])", R"(["A", "B", "C", "F"])", "[1, 2, 3, 6]", + R"(["a", "b", "c", "f"])", R"(["A", "B", "C", "F"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + expected_table.push_back(expected_result); + + expected_result_string = {R"(["j", "l", "n"])", R"(["J", "L", "N"])", "[8, 10, 12]", + R"(["j", "l", "n"])", R"(["J", "L", "N"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + expected_table.push_back(expected_result); + + std::vector>> expected_table_vec; + { + std::vector> iterator_expected_table; + std::shared_ptr expected_result; + std::vector expected_result_string = { + R"(["a", "b"])", R"(["A", "B"])", "[1, 2]", R"(["a", "b"])", R"(["A", "B"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + iterator_expected_table.push_back(expected_result); + expected_result_string = {R"(["c", "f"])", R"(["C", "F"])", "[3, 6]", R"(["c", "f"])", + R"(["C", "F"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + iterator_expected_table.push_back(expected_result); + expected_table_vec.push_back(iterator_expected_table); + } + { + std::vector> iterator_expected_table; + std::shared_ptr expected_result; + std::vector expected_result_string = { + R"(["j", "l"])", R"(["J", "L"])", "[8, 10]", R"(["j", "l"])", R"(["J", "L"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + iterator_expected_table.push_back(expected_result); + expected_result_string = {R"(["n"])", R"(["N"])", "[12]", R"(["n"])", R"(["N"])"}; + MakeInputBatch(expected_result_string, schema_table, &expected_result); + iterator_expected_table.push_back(expected_result); + expected_table_vec.push_back(iterator_expected_table); + } + + ////////////////////// evaluate ////////////////////// + for (auto batch : table_0) { + ASSERT_NOT_OK(expr_build->evaluate(batch, &dummy_result_batches)); + } + std::shared_ptr build_result_iterator; + std::shared_ptr probe_result_iterator_base; + ASSERT_NOT_OK(expr_build->finish(&build_result_iterator)); + ASSERT_NOT_OK(expr_probe->finish(&probe_result_iterator_base)); + + auto probe_result_iterator = + std::dynamic_pointer_cast>( + probe_result_iterator_base); + probe_result_iterator->SetDependencies({build_result_iterator}); + + for (int i = 0; i < 2; i++) { + auto right_batch = table_1[i]; + + std::shared_ptr result_batch; + std::vector> input; + for (int i = 0; i < right_batch->num_columns(); i++) { + input.push_back(right_batch->column(i)); + } + + ASSERT_NOT_OK(probe_result_iterator->Process(input, &result_batch)); + ASSERT_NOT_OK(Equals(*(expected_table[i]).get(), *result_batch.get())); + + std::shared_ptr iterator_result_batch; + int j = 0; + while (probe_result_iterator->HasNext()) { + ASSERT_NOT_OK(probe_result_iterator->Next(&iterator_result_batch)); + ASSERT_NOT_OK( + Equals(*(expected_table_vec[i][j]).get(), *iterator_result_batch.get())); + j++; + } + } +} + TEST(TestArrowComputeWSCG, JoinWOCGTestInnerJoinNaN) { ////////////////////// prepare expr_vector /////////////////////// auto table0_f0 = field("table0_f0", float64());