Fix 'physical explain' of merge limit

Signed-off-by: Jin Hai <[email protected]>
infiniflow · Dec 15, 2024 · 0b1ce60 · 0b1ce60
1 parent 5ecf1f5
commit 0b1ce60
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 8 deletions.
diff --git a/python/test_pysdk/test_py_import.py b/python/test_pysdk/test_py_import.py
@@ -420,6 +420,12 @@ def test_import_with_different_size(self, check_data, data_size, suffix):
                                         {"c1": {"type": "int"}, "c2": {"type": "varchar"}})
 
         test_csv_dir = common_values.TEST_TMP_DIR + "pysdk_test_import_with_different_size.csv"
+
+        with open(test_csv_dir, 'r') as file:
+            row_count = sum(1 for line in file)
+
+        assert row_count == data_size
+
         res = table_obj.import_data(test_csv_dir)
         assert res.error_code == ErrorCode.OK
 

diff --git a/src/executor/explain_physical_plan.cpp b/src/executor/explain_physical_plan.cpp
@@ -239,18 +239,22 @@ void ExplainPhysicalPlan::Explain(const PhysicalOperator *op, SharedPtr<Vector<S
             break;
         }
         case PhysicalOperatorType::kSource: {
+            Explain((PhysicalSource *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kSink: {
+            Explain((PhysicalSink *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kInvalid: {
             break;
         }
         case PhysicalOperatorType::kParallelAggregate: {
+            Explain((PhysicalParallelAggregate *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kMergeParallelAggregate: {
+            Explain((PhysicalMergeParallelAggregate *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kIntersect: {
@@ -266,13 +270,15 @@ void ExplainPhysicalPlan::Explain(const PhysicalOperator *op, SharedPtr<Vector<S
             break;
         }
         case PhysicalOperatorType::kMergeLimit: {
+            Explain((PhysicalMergeLimit *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kMergeTop: {
             Explain((PhysicalMergeTop *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kMergeSort: {
+            Explain((PhysicalMergeSort *)op, result, intent_size);
             break;
         }
         case PhysicalOperatorType::kMergeKnn: {
@@ -2282,6 +2288,30 @@ void ExplainPhysicalPlan::Explain(const PhysicalMergeLimit *merge_limit_node,
     }
     explain_header_str += "(" + std::to_string(merge_limit_node->node_id()) + ")";
     result->emplace_back(MakeShared<String>(explain_header_str));
+
+    {
+        String limit_value_str = String(intent_size, ' ') + " - limit: ";
+        ExplainLogicalPlan::Explain(merge_limit_node->limit_expr().get(), limit_value_str);
+        result->emplace_back(MakeShared<String>(limit_value_str));
+    }
+
+    if (merge_limit_node->offset_expr().get() != 0) {
+        String offset_value_str = String(intent_size, ' ') + " - offset: ";
+        ExplainLogicalPlan::Explain(merge_limit_node->offset_expr().get(), offset_value_str);
+        result->emplace_back(MakeShared<String>(offset_value_str));
+    }
+
+    // Output column
+    {
+        String output_columns_str = String(intent_size, ' ') + " - output columns: [";
+        SharedPtr<Vector<String>> output_columns = merge_limit_node->GetOutputNames();
+        SizeT column_count = output_columns->size();
+        for (SizeT idx = 0; idx < column_count - 1; ++idx) {
+            output_columns_str += output_columns->at(idx) + ", ";
+        }
+        output_columns_str += output_columns->back() + "]";
+        result->emplace_back(MakeShared<String>(output_columns_str));
+    }
 }
 
 void ExplainPhysicalPlan::Explain(const PhysicalMergeTop *merge_top_node, SharedPtr<Vector<SharedPtr<String>>> &result, i64 intent_size) {

diff --git a/src/executor/operator/physical_merge_limit.cppm b/src/executor/operator/physical_merge_limit.cppm
@@ -56,6 +56,10 @@ public:
         return 0;
     }
 
+    [[nodiscard]] inline const SharedPtr<BaseExpression> &limit_expr() const { return limit_expr_; }
+
+    [[nodiscard]] inline const SharedPtr<BaseExpression> &offset_expr() const { return offset_expr_; }
+
 private:
     SharedPtr<BaseExpression> limit_expr_{};
     SharedPtr<BaseExpression> offset_expr_{};

diff --git a/src/executor/physical_planner.cpp b/src/executor/physical_planner.cpp
@@ -747,29 +747,30 @@ UniquePtr<PhysicalOperator> PhysicalPlanner::BuildLimit(const SharedPtr<LogicalN
 
     SharedPtr<LogicalLimit> logical_limit = static_pointer_cast<LogicalLimit>(logical_operator);
     UniquePtr<PhysicalOperator> input_physical_operator = BuildPhysicalOperator(input_logical_node);
-    if (input_physical_operator->TaskletCount() <= 1) {
+
+    i64 offset = 0;
+    if (logical_limit->offset_expression_.get() != nullptr) {
+        offset = (static_pointer_cast<ValueExpression>(logical_limit->offset_expression_))->GetValue().value_.big_int;
+    }
+
+    if (input_physical_operator->TaskletCount() <= 1 or offset != 0) {
         return MakeUnique<PhysicalLimit>(logical_operator->node_id(),
                                          std::move(input_physical_operator),
                                          logical_limit->limit_expression_,
                                          logical_limit->offset_expression_,
                                          logical_operator->load_metas(),
                                          logical_limit->total_hits_count_flag_);
     } else {
-        i64 child_limit = (static_pointer_cast<ValueExpression>(logical_limit->limit_expression_))->GetValue().value_.big_int;
-
-        if (logical_limit->offset_expression_.get() != nullptr) {
-            child_limit += (static_pointer_cast<ValueExpression>(logical_limit->offset_expression_))->GetValue().value_.big_int;
-        }
         auto child_limit_op = MakeUnique<PhysicalLimit>(logical_operator->node_id(),
                                                         std::move(input_physical_operator),
-                                                        MakeShared<ValueExpression>(Value::MakeBigInt(child_limit)),
+                                                        logical_limit->limit_expression_,
                                                         nullptr,
                                                         logical_operator->load_metas(),
                                                         logical_limit->total_hits_count_flag_);
         return MakeUnique<PhysicalMergeLimit>(query_context_ptr_->GetNextNodeID(),
                                               std::move(child_limit_op),
                                               logical_limit->limit_expression_,
-                                              logical_limit->offset_expression_,
+                                              nullptr,
                                               MakeShared<Vector<LoadMeta>>());
     }
 }