Merge branch 'main' of https://github.com/datafuselabs/databend into add_runtime_bloom_filter_for_merge_into

Author: JackTan25

src/common/metrics/src/metrics/session.rs

@@ -13,18 +13,31 @@
 // limitations under the License.
 
 use std::sync::LazyLock;
+use std::time::Duration;
 
 use crate::register_counter;
 use crate::register_gauge;
+use crate::register_histogram_in_milliseconds;
 use crate::Counter;
 use crate::Gauge;
+use crate::Histogram;
 
 pub static SESSION_CONNECT_NUMBERS: LazyLock<Counter> =
     LazyLock::new(|| register_counter("session_connect_numbers"));
 pub static SESSION_CLOSE_NUMBERS: LazyLock<Counter> =
     LazyLock::new(|| register_counter("session_close_numbers"));
 pub static SESSION_ACTIVE_CONNECTIONS: LazyLock<Gauge> =
     LazyLock::new(|| register_gauge("session_connections"));
+pub static SESSION_QUQUED_QUERIES: LazyLock<Gauge> =
+    LazyLock::new(|| register_gauge("session_queued_queries"));
+pub static SESSION_QUEUE_ABORT_COUNT: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("session_queue_abort_count"));
+pub static SESSION_QUEUE_ACQUIRE_ERROR_COUNT: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("session_queue_acquire_error_count"));
+pub static SESSION_QUEUE_ACQUIRE_TIMEOUT_COUNT: LazyLock<Counter> =
+    LazyLock::new(|| register_counter("session_queue_acquire_timeout_count"));
+pub static SESSION_QUEUE_ACQUIRE_DURATION_MS: LazyLock<Histogram> =
+    LazyLock::new(|| register_histogram_in_milliseconds("session_queue_acquire_duration_ms"));
 
 pub fn incr_session_connect_numbers() {
     SESSION_CONNECT_NUMBERS.inc();
@@ -37,3 +50,23 @@ pub fn incr_session_close_numbers() {
 pub fn set_session_active_connections(num: usize) {
     SESSION_ACTIVE_CONNECTIONS.set(num as i64);
 }
+
+pub fn set_session_queued_queries(num: usize) {
+    SESSION_QUQUED_QUERIES.set(num as i64);
+}
+
+pub fn incr_session_queue_abort_count() {
+    SESSION_QUEUE_ABORT_COUNT.inc();
+}
+
+pub fn incr_session_queue_acquire_error_count() {
+    SESSION_QUEUE_ACQUIRE_ERROR_COUNT.inc();
+}
+
+pub fn incr_session_queue_acquire_timeout_count() {
+    SESSION_QUEUE_ACQUIRE_TIMEOUT_COUNT.inc();
+}
+
+pub fn record_session_queue_acquire_duration_ms(duration: Duration) {
+    SESSION_QUEUE_ACQUIRE_DURATION_MS.observe(duration.as_millis() as f64);
+}

src/query/ast/src/parser/statement.rs

@@ -1986,7 +1986,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
     );
 
-    let begin = value(Statement::Begin, rule! { BEGIN });
+    let begin = value(Statement::Begin, rule! { BEGIN ~ (TRANSACTION)? });
     let commit = value(Statement::Commit, rule! { COMMIT });
     let abort = value(Statement::Abort, rule! { ABORT | ROLLBACK });
 

src/query/ast/src/parser/token.rs

@@ -1213,6 +1213,8 @@ pub enum TokenKind {
     UNTIL,
     #[token("BEGIN", ignore(ascii_case))]
     BEGIN,
+    #[token("TRANSACTION", ignore(ascii_case))]
+    TRANSACTION,
     #[token("COMMIT", ignore(ascii_case))]
     COMMIT,
     #[token("ABORT", ignore(ascii_case))]
@@ -1541,6 +1543,7 @@ impl TokenKind {
             // | TokenKind::TABLESAMPLE
             | TokenKind::THEN
             | TokenKind::TRAILING
+            | TokenKind::TRANSACTION
             | TokenKind::TRUE
             // | TokenKind::UNIQUE
             //| TokenKind::USER

src/query/service/src/sessions/queue_mgr.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::fmt::Display;
 use std::future::Future;
 use std::hash::Hash;
 use std::pin::Pin;
@@ -30,6 +31,11 @@ use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_metrics::session::incr_session_queue_abort_count;
+use databend_common_metrics::session::incr_session_queue_acquire_error_count;
+use databend_common_metrics::session::incr_session_queue_acquire_timeout_count;
+use databend_common_metrics::session::record_session_queue_acquire_duration_ms;
+use databend_common_metrics::session::set_session_queued_queries;
 use log::info;
 use parking_lot::Mutex;
 use pin_project_lite::pin_project;
@@ -41,7 +47,7 @@ use tokio::time::error::Elapsed;
 use crate::sessions::QueryContext;
 
 pub trait QueueData: Send + Sync + 'static {
-    type Key: Send + Sync + Eq + Hash + Clone + 'static;
+    type Key: Send + Sync + Eq + Hash + Display + Clone + 'static;
 
     fn get_key(&self) -> Self::Key;
 
@@ -91,12 +97,14 @@ impl<Data: QueueData> QueueManager<Data> {
     pub fn remove(&self, key: Data::Key) -> bool {
         let mut queue = self.queue.lock();
         if let Some(inner) = queue.remove(&key) {
+            set_session_queued_queries(queue.len());
             inner.waker.wake();
             inner.is_abort.store(true, Ordering::SeqCst);
-            return true;
+            true
+        } else {
+            set_session_queued_queries(queue.len());
+            false
         }
-
-        false
     }
 
     pub async fn acquire(self: &Arc<Self>, data: Data) -> Result<AcquireQueueGuard> {
@@ -106,20 +114,43 @@ impl<Data: QueueData> QueueManager<Data> {
             tokio::time::timeout(timeout, self.semaphore.clone().acquire_owned()),
             self.clone(),
         );
+        let start_time = SystemTime::now();
 
-        future.await
+        match future.await {
+            Ok(v) => {
+                record_session_queue_acquire_duration_ms(start_time.elapsed().unwrap_or_default());
+                Ok(v)
+            }
+            Err(e) => {
+                match e.code() {
+                    ErrorCode::ABORTED_QUERY => {
+                        incr_session_queue_abort_count();
+                    }
+                    ErrorCode::TIMEOUT => {
+                        incr_session_queue_acquire_timeout_count();
+                    }
+                    _ => {
+                        incr_session_queue_acquire_error_count();
+                    }
+                }
+                Err(e)
+            }
+        }
     }
 
     pub(crate) fn add_entity(&self, inner: Inner<Data>) -> Data::Key {
         let key = inner.data.get_key();
         let mut queue = self.queue.lock();
         queue.insert(key.clone(), inner);
+        set_session_queued_queries(queue.len());
         key
     }
 
     pub(crate) fn remove_entity(&self, key: &Data::Key) -> Option<Arc<Data>> {
         let mut queue = self.queue.lock();
-        queue.remove(key).map(|inner| inner.data.clone())
+        let data = queue.remove(key).map(|inner| inner.data.clone());
+        set_session_queued_queries(queue.len());
+        data
     }
 }
 

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -283,17 +283,14 @@ pub fn optimize_query(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<SE
         enable_distributed_query,
     )?;
 
+    if opt_ctx.enable_join_reorder {
+        s_expr =
+            RecursiveOptimizer::new([RuleID::CommuteJoin].as_slice(), &opt_ctx).run(&s_expr)?;
+    }
+
     // Cascades optimizer may fail due to timeout, fallback to heuristic optimizer in this case.
     s_expr = match cascades.optimize(s_expr.clone()) {
         Ok(mut s_expr) => {
-            let rules = if opt_ctx.enable_join_reorder {
-                [RuleID::EliminateEvalScalar, RuleID::CommuteJoin].as_slice()
-            } else {
-                [RuleID::EliminateEvalScalar].as_slice()
-            };
-
-            s_expr = RecursiveOptimizer::new(rules, &opt_ctx).run(&s_expr)?;
-
             // Push down sort and limit
             // TODO(leiysky): do this optimization in cascades optimizer
             if enable_distributed_query {
@@ -308,10 +305,6 @@ pub fn optimize_query(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<SE
                 "CascadesOptimizer failed, fallback to heuristic optimizer: {}",
                 e
             );
-
-            s_expr =
-                RecursiveOptimizer::new(&[RuleID::EliminateEvalScalar], &opt_ctx).run(&s_expr)?;
-
             if enable_distributed_query {
                 s_expr = optimize_distributed_query(opt_ctx.table_ctx.clone(), &s_expr)?;
             }
@@ -320,6 +313,9 @@ pub fn optimize_query(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<SE
         }
     };
 
+    s_expr =
+        RecursiveOptimizer::new([RuleID::EliminateEvalScalar].as_slice(), &opt_ctx).run(&s_expr)?;
+
     Ok(s_expr)
 }
 

tests/sqllogictests/suites/base/14_transaction/14_0003_merge_into.test

@@ -69,5 +69,80 @@ SELECT employee_id, salary FROM salaries order by employee_id;
 4   55000.00
 
 
+statement ok
+drop database test_txn_merge_into;
+
+
+statement ok
+create or replace database test_txn_merge_into;
+
+statement ok
+use test_txn_merge_into;
+
+statement ok
+set enable_experimental_merge_into = 1;
+
+statement ok
+CREATE TABLE employees (
+    employee_id INT,
+    employee_name VARCHAR(255),
+    department VARCHAR(255)
+);
+
+statement ok
+CREATE TABLE salaries (
+    employee_id INT,
+    salary DECIMAL(10, 2)
+);
+
+statement ok
+BEGIN TRANSACTION;
+
+statement ok
+INSERT INTO employees VALUES
+    (1, 'Alice', 'HR'),
+    (2, 'Bob', 'IT'),
+    (3, 'Charlie', 'Finance'),
+    (4, 'David', 'HR');
+
+statement ok
+INSERT INTO salaries VALUES
+    (1, 50000.00),
+    (2, 60000.00);
+
+statement ok
+MERGE INTO salaries
+    USING (SELECT * FROM employees) AS employees
+    ON salaries.employee_id = employees.employee_id
+    WHEN MATCHED AND employees.department = 'HR' THEN
+        UPDATE SET
+            salaries.salary = salaries.salary + 1000.00
+    WHEN MATCHED THEN
+        UPDATE SET
+            salaries.salary = salaries.salary + 500.00
+    WHEN NOT MATCHED THEN
+        INSERT (employee_id, salary)
+            VALUES (employees.employee_id, 55000.00);
+
+query IF
+SELECT employee_id, salary FROM salaries order by employee_id;
+----
+1   51000.00
+2   60500.00
+3   55000.00
+4   55000.00
+
+statement ok
+COMMIT;
+
+query IF
+SELECT employee_id, salary FROM salaries order by employee_id;
+----
+1   51000.00
+2   60500.00
+3   55000.00
+4   55000.00
+
+
 statement ok
 drop database test_txn_merge_into;