fix(query): fix bind internal column (#17463)

* fix(query): fix bind internal column

* fix

* fix merge into

Author: b41sh

src/query/sql/src/planner/binder/bind_context.rs

@@ -14,6 +14,7 @@
 
 use std::collections::btree_map;
 use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::hash::Hash;
@@ -147,8 +148,8 @@ pub struct BindContext {
 
     pub columns: Vec<ColumnBinding>,
 
-    // map internal column id to (table_index, column_index)
-    pub bound_internal_columns: BTreeMap<ColumnId, (IndexType, IndexType)>,
+    // map internal column: (table_index, column_id) -> column_index
+    pub bound_internal_columns: BTreeMap<(IndexType, ColumnId), IndexType>,
 
     pub aggregate_info: AggregateInfo,
 
@@ -578,35 +579,49 @@ impl BindContext {
     }
 
     fn get_internal_column_table_index(
+        &self,
         column_binding: &InternalColumnBinding,
-        metadata: MetadataRef,
     ) -> Result<IndexType> {
-        let metadata = metadata.read();
-
         if let Some(table_name) = &column_binding.table_name {
-            metadata
-                .get_table_index(column_binding.database_name.as_deref(), table_name)
-                .ok_or_else(|| {
-                    ErrorCode::TableInfoError(format!(
-                        "Table `{table_name}` is not found in the metadata"
-                    ))
-                })
+            for column in &self.columns {
+                if column_binding.database_name.is_some()
+                    && column.database_name != column_binding.database_name
+                {
+                    continue;
+                }
+                if column.table_name != column_binding.table_name {
+                    continue;
+                }
+                if let Some(table_index) = column.table_index {
+                    return Ok(table_index);
+                }
+            }
+            Err(ErrorCode::SemanticError(format!(
+                "Table `{table_name}` is not found in the bind context"
+            )))
         } else {
-            let tables = metadata
-                .tables()
-                .iter()
-                .filter(|t| !t.is_source_of_index())
-                .collect::<Vec<_>>();
-            debug_assert!(!tables.is_empty());
-
-            if tables.len() > 1 {
+            let mut table_indices = BTreeSet::new();
+            for column in &self.columns {
+                if column.table_name.is_none() {
+                    continue;
+                }
+                if let Some(table_index) = column.table_index {
+                    table_indices.insert(table_index);
+                }
+            }
+            if table_indices.is_empty() {
+                return Err(ErrorCode::SemanticError(format!(
+                    "The table of the internal column `{}` is not found",
+                    column_binding.internal_column.column_name()
+                )));
+            }
+            if table_indices.len() > 1 {
                 return Err(ErrorCode::SemanticError(format!(
                     "The table of the internal column `{}` is ambiguous",
                     column_binding.internal_column.column_name()
                 )));
             }
-
-            Ok(tables[0].index())
+            Ok(*table_indices.first().unwrap())
         }
     }
 
@@ -615,25 +630,31 @@ impl BindContext {
         &mut self,
         column_binding: &InternalColumnBinding,
         metadata: MetadataRef,
+        table_index: Option<IndexType>,
         visible: bool,
     ) -> Result<ColumnBinding> {
         let column_id = column_binding.internal_column.column_id();
-        let (table_index, column_index, new) = match self.bound_internal_columns.entry(column_id) {
-            btree_map::Entry::Vacant(e) => {
-                let table_index =
-                    BindContext::get_internal_column_table_index(column_binding, metadata.clone())?;
-                let mut metadata = metadata.write();
-                let column_index = metadata
-                    .add_internal_column(table_index, column_binding.internal_column.clone());
-                e.insert((table_index, column_index));
-                (table_index, column_index, true)
-            }
-            btree_map::Entry::Occupied(e) => {
-                let (table_index, column_index) = e.get();
-                (*table_index, *column_index, false)
-            }
+        let table_index = if let Some(table_index) = table_index {
+            table_index
+        } else {
+            self.get_internal_column_table_index(column_binding)?
         };
 
+        let (column_index, is_new) =
+            match self.bound_internal_columns.entry((table_index, column_id)) {
+                btree_map::Entry::Vacant(e) => {
+                    let mut metadata = metadata.write();
+                    let column_index = metadata
+                        .add_internal_column(table_index, column_binding.internal_column.clone());
+                    e.insert(column_index);
+                    (column_index, true)
+                }
+                btree_map::Entry::Occupied(e) => {
+                    let column_index = e.get();
+                    (*column_index, false)
+                }
+            };
+
         let metadata = metadata.read();
         let table = metadata.table(table_index);
         if !table.table().supported_internal_column(column_id) {
@@ -660,7 +681,7 @@ impl BindContext {
         .table_index(Some(table_index))
         .build();
 
-        if new {
+        if is_new {
             debug_assert!(!self.columns.iter().any(|c| c == &column_binding));
             self.columns.push(column_binding.clone());
         }

src/query/sql/src/planner/binder/bind_mutation/mutation_expression.rs

@@ -417,6 +417,7 @@ impl Binder {
         let column_binding = match bind_context.add_internal_column_binding(
             &row_id_column_binding,
             self.metadata.clone(),
+            Some(table_index),
             true,
         ) {
             Ok(column_binding) => column_binding,

src/query/sql/src/planner/binder/binder.rs

@@ -975,7 +975,7 @@ impl<'a> Binder {
 
         let mut has_score = false;
         let mut has_matched = false;
-        for column_id in bound_internal_columns.keys() {
+        for (_, column_id) in bound_internal_columns.keys() {
             if *column_id == SEARCH_SCORE_COLUMN_ID {
                 has_score = true;
             } else if *column_id == SEARCH_MATCHED_COLUMN_ID {
@@ -988,7 +988,7 @@ impl<'a> Binder {
             ));
         }
 
-        for (table_index, column_index) in bound_internal_columns.values() {
+        for ((table_index, _), column_index) in bound_internal_columns.iter() {
             let inverted_index = inverted_index_map.shift_remove(table_index).map(|mut i| {
                 i.has_score = has_score;
                 i

src/query/sql/src/planner/semantic/type_check.rs

@@ -295,6 +295,7 @@ impl<'a> TypeChecker<'a> {
                         let column = self.bind_context.add_internal_column_binding(
                             &column,
                             self.metadata.clone(),
+                            None,
                             true,
                         )?;
                         let data_type = *column.data_type.clone();
@@ -2199,6 +2200,7 @@ impl<'a> TypeChecker<'a> {
         let column = self.bind_context.add_internal_column_binding(
             &internal_column_binding,
             self.metadata.clone(),
+            None,
             false,
         )?;
 
@@ -2630,6 +2632,7 @@ impl<'a> TypeChecker<'a> {
         let column = self.bind_context.add_internal_column_binding(
             &internal_column_binding,
             self.metadata.clone(),
+            None,
             false,
         )?;
 

tests/sqllogictests/suites/base/05_ddl/05_0031_internal_column.test

@@ -57,11 +57,21 @@ SELECT a,`05_0031_t_1`._row_id FROM `05_0031_t` NATURAL JOIN `05_0031_t_1` order
 3 18446739675663040512
 3 18446739675663040513
 
+query III
+SELECT a, `05_0031_t`._row_id, `05_0031_t_1`._row_id FROM `05_0031_t`, `05_0031_t_1` order by a, `05_0031_t`._row_id;
+----
+1 18446735277616529408 18446739675663040512
+1 18446735277616529408 18446739675663040513
+2 18446735277616529409 18446739675663040512
+2 18446735277616529409 18446739675663040513
+3 18446739675663040512 18446739675663040512
+3 18446739675663040512 18446739675663040513
+
 # issue #11126
 # ===============
 
 statement ok
-create table t_11126(a Int64)
+create or replace table t_11126(a Int64)
 
 statement ok
 insert into t_11126 select * from numbers(100);
@@ -75,6 +85,26 @@ select a, _row_id from t_11126 order by a desc limit 5;
 96 18446739675663040608
 95 18446739675663040607
 
+# issue 17327
+
+query IF
+with tt as (select sum(a) as a from t_11126) select t_11126._row_id, t_11126.a/tt.a from t_11126, tt limit 5;
+----
+18446739675663040512 0.0
+18446739675663040513 0.00020202020202020202
+18446739675663040514 0.00040404040404040404
+18446739675663040515 0.0006060606060606061
+18446739675663040516 0.0008080808080808081
+
+query IF
+with tt as (select sum(a) as a from t_11126) select t_11126._row_id, t_11126.a/tt.a from tt, t_11126 limit 5;
+----
+18446739675663040512 0.0
+18446739675663040513 0.00020202020202020202
+18446739675663040514 0.00040404040404040404
+18446739675663040515 0.0006060606060606061
+18446739675663040516 0.0008080808080808081
+
 statement ok
 drop table t_11126
 
@@ -104,4 +134,4 @@ statement ok
 drop table t_11772_1;
 
 statement ok
-drop table t_11772_2;
+drop table t_11772_2;