From 39f8b173744cc1d476c36705367c1649d9277d92 Mon Sep 17 00:00:00 2001
From: Kould <kould2333@gmail.com>
Date: Fri, 29 Dec 2023 18:31:58 +0800
Subject: [PATCH] fix: remapper error when Project has no first column or
 column order is reverse order (#398)

---
 src/planner/bound/base_table_ref.cppm | 17 ++++++++++
 src/planner/optimizer/lazy_load.cpp   | 12 +++----
 test/sql/dql/select.slt               | 48 +++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/src/planner/bound/base_table_ref.cppm b/src/planner/bound/base_table_ref.cppm
index ef4be06ed2..8166cc8d35 100644
--- a/src/planner/bound/base_table_ref.cppm
+++ b/src/planner/bound/base_table_ref.cppm
@@ -53,6 +53,23 @@ public:
         replace_field<SharedPtr<DataType>>(*column_types_, indices);
     };
 
+    void RetainColumnByIds(Vector<SizeT> &&ids) {
+        if (ids.empty()) {
+            return;
+        }
+        Vector<SizeT> indices;
+        indices.reserve(ids.size());
+
+        std::sort(ids.begin(), ids.end());
+        for (SizeT i = 0, ids_i = 0; i < column_ids_.size() && ids_i < ids.size(); ++i) {
+            if (column_ids_[i] == ids[ids_i]) {
+                indices.push_back(i);
+                ids_i ++;
+            }
+        }
+        RetainColumnByIndices(Move(indices));
+    }
+
     SharedPtr<String> schema_name() const { return table_entry_ptr_->table_collection_meta_->db_entry_->db_name_; }
 
     SharedPtr<String> table_name() const { return table_entry_ptr_->table_collection_name_; }
diff --git a/src/planner/optimizer/lazy_load.cpp b/src/planner/optimizer/lazy_load.cpp
index 33a1a976c0..f1d9fc0dd7 100644
--- a/src/planner/optimizer/lazy_load.cpp
+++ b/src/planner/optimizer/lazy_load.cpp
@@ -112,27 +112,27 @@ void CleanScan::VisitNode(LogicalNode &op) {
     switch (op.operator_type()) {
         case LogicalNodeType::kTableScan: {
             auto table_scan = dynamic_cast<LogicalTableScan &>(op);
-            Vector<SizeT> project_indices = LoadedColumn(last_op_load_metas_.get(), table_scan.base_table_ref_.get());
+            Vector<SizeT> project_ids = LoadedColumn(last_op_load_metas_.get(), table_scan.base_table_ref_.get());
 
             scan_table_indexes_.push_back(table_scan.base_table_ref_->table_index_);
-            table_scan.base_table_ref_->RetainColumnByIndices(Move(project_indices));
+            table_scan.base_table_ref_->RetainColumnByIds(Move(project_ids));
             table_scan.add_row_id_ = true;
             break;
         }
         case LogicalNodeType::kKnnScan: {
             auto knn_scan = dynamic_cast<LogicalKnnScan &>(op);
-            Vector<SizeT> project_indices = LoadedColumn(last_op_load_metas_.get(), knn_scan.base_table_ref_.get());
+            Vector<SizeT> project_ids = LoadedColumn(last_op_load_metas_.get(), knn_scan.base_table_ref_.get());
 
             scan_table_indexes_.push_back(knn_scan.base_table_ref_->table_index_);
-            knn_scan.base_table_ref_->RetainColumnByIndices(Move(project_indices));
+            knn_scan.base_table_ref_->RetainColumnByIds(Move(project_ids));
             break;
         }
         case LogicalNodeType::kMatch: {
             auto match = dynamic_cast<LogicalMatch &>(op);
-            Vector<SizeT> project_indices = LoadedColumn(last_op_load_metas_.get(), match.base_table_ref_.get());
+            Vector<SizeT> project_ids = LoadedColumn(last_op_load_metas_.get(), match.base_table_ref_.get());
 
             scan_table_indexes_.push_back(match.base_table_ref_->table_index_);
-            match.base_table_ref_->RetainColumnByIndices(Move(project_indices));
+            match.base_table_ref_->RetainColumnByIds(Move(project_ids));
             break;
         }
         case LogicalNodeType::kLimit:
diff --git a/test/sql/dql/select.slt b/test/sql/dql/select.slt
index 77d12d5799..8d0444bb9e 100644
--- a/test/sql/dql/select.slt
+++ b/test/sql/dql/select.slt
@@ -7,17 +7,65 @@ CREATE TABLE select1 (id INTEGER PRIMARY KEY, name VARCHAR, age INTEGER);
 statement ok
 CREATE TABLE select2 (id INTEGER , age INTEGER);
 
+statement ok
+CREATE TABLE select3 (c1 INTEGER, c2 INTEGER, c3 INTEGER);
+
 # copy data from csv file
 query I
 COPY select2 FROM '/tmp/infinity/test_data/nation.csv' WITH ( DELIMITER ',' );
 ----
 
+statement ok
+INSERT INTO select3 VALUES(0,1,2),(3,4,5),(6,7,8);
+
 #query ITI
 #SELECT * FROM select1 ORDER by age ASC;
 #----
 #2 Jane 25
 #1 John 30
 
+query III
+SELECT * from select3;
+----
+0 1 2
+3 4 5
+6 7 8
+
+query III
+SELECT c1, c2, c3 from select3;
+----
+0 1 2
+3 4 5
+6 7 8
+
+query III
+SELECT c2, c3 from select3;
+----
+1 2
+4 5
+7 8
+
+query III
+SELECT c2 from select3;
+----
+1
+4
+7
+
+query III
+SELECT c3, c2, c1 from select3;
+----
+2 1 0
+5 4 3
+8 7 6
+
+query III
+SELECT c2, c1 from select3;
+----
+1 0
+4 3
+7 6
+
 statement ok
 DROP TABLE select1;