diff --git a/Cargo.lock b/Cargo.lock index ba758429c41fe..e4ec6dae31c7e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5052,6 +5052,7 @@ dependencies = [ "databend-common-sql", "databend-common-statistics", "databend-common-storage", + "databend-common-storages-parquet", "databend-common-tracing", "databend-common-users", "databend-enterprise-fail-safe", @@ -6642,6 +6643,7 @@ dependencies = [ "databend-common-catalog", "databend-common-config", "databend-common-exception", + "databend-common-expression", "databend-common-metrics", "databend-storages-common-index", "databend-storages-common-table-meta", diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs index a8d494394a32e..3e7f549d84f7b 100644 --- a/src/query/service/src/interpreters/access/privilege_access.rs +++ b/src/query/service/src/interpreters/access/privilege_access.rs @@ -756,6 +756,16 @@ impl PrivilegeAccess { ) -> Result<()> { let session = self.ctx.get_current_session(); + // Direct user grants do not require resolving an ownership object. + if self + .ctx + .get_current_user()? + .grants + .verify_privilege(grant_object, privilege) + { + return Ok(()); + } + let verify_ownership = match grant_object { GrantObject::Database(_, _) | GrantObject::Table(_, _, _) diff --git a/src/query/service/src/interpreters/common/log.rs b/src/query/service/src/interpreters/common/log.rs index 0b5b55112187d..b3a1bb2913d26 100644 --- a/src/query/service/src/interpreters/common/log.rs +++ b/src/query/service/src/interpreters/common/log.rs @@ -16,11 +16,14 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::time::SystemTime; +use databend_common_base::runtime::ThreadTracker; use databend_common_base::runtime::profile::ProfileDesc; use databend_common_base::runtime::profile::ProfileStatisticsName; use databend_common_base::runtime::profile::get_statistics_desc; +use databend_common_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_pipeline::core::PlanProfile; +use databend_common_tracing::HistoryConfig; use log::error; use log::info; @@ -33,6 +36,43 @@ use crate::sessions::TableContextQueryIdentity; use crate::sessions::TableContextQueryProfile; use crate::sessions::TableContextSpillProgress; +const QUERY_HISTORY_TABLE: &str = "query_history"; +const PROFILE_HISTORY_TABLE: &str = "profile_history"; +const ACCESS_HISTORY_TABLE: &str = "access_history"; + +fn history_table_log_enabled(history: &HistoryConfig, history_table_name: &str) -> bool { + history.on + && history + .tables + .iter() + .any(|table| table.table_name.eq_ignore_ascii_case(history_table_name)) +} + +fn captured_info_log_enabled() -> bool { + ThreadTracker::capture_log_settings().is_some_and(|settings| { + settings.queue.is_some() && settings.level >= log::LevelFilter::Info + }) +} + +pub(crate) fn query_log_enabled() -> bool { + let log = &GlobalConfig::instance().log; + log.query.on + || history_table_log_enabled(&log.history, QUERY_HISTORY_TABLE) + || captured_info_log_enabled() +} + +pub(crate) fn profile_log_enabled() -> bool { + let log = &GlobalConfig::instance().log; + log.profile.on + || history_table_log_enabled(&log.history, PROFILE_HISTORY_TABLE) + || captured_info_log_enabled() +} + +pub(crate) fn access_log_enabled() -> bool { + let log = &GlobalConfig::instance().log; + history_table_log_enabled(&log.history, ACCESS_HISTORY_TABLE) || captured_info_log_enabled() +} + pub fn log_query_start(ctx: &QueryContext) { InterpreterMetrics::record_query_start(ctx); let now = SystemTime::now(); @@ -42,7 +82,9 @@ pub fn log_query_start(ctx: &QueryContext) { SessionManager::instance().status.write().query_start(now); } - if let Err(error) = InterpreterQueryLog::log_start(ctx, now, None) { + if query_log_enabled() + && let Err(error) = InterpreterQueryLog::log_start(ctx, now, None) + { error!("Failed to log query start: {:?}", error) } } @@ -52,6 +94,7 @@ pub fn log_query_finished(ctx: &QueryContext, error: Option) { InterpreterMetrics::record_query_finished(ctx, error.clone()); let now = SystemTime::now(); + ctx.set_finish_time(now); let session = ctx.get_current_session(); session.get_status().write().query_finish(); @@ -72,11 +115,16 @@ pub fn log_query_finished(ctx: &QueryContext, error: Option) { info!(memory:? = ctx.get_node_peek_memory_usage(); "total memory usage"); - // databend::log::profile - let query_profiles = ctx.get_query_profiles(); + let query_log_is_enabled = query_log_enabled(); + let profile_log_is_enabled = profile_log_enabled(); + let query_profiles = if query_log_is_enabled || profile_log_is_enabled { + ctx.get_query_profiles() + } else { + Vec::new() + }; let has_profiles = !query_profiles.is_empty(); - if has_profiles { + if has_profiles && profile_log_is_enabled { #[derive(serde::Serialize)] struct QueryProfiles { query_id: String, @@ -98,10 +146,62 @@ pub fn log_query_finished(ctx: &QueryContext, error: Option) { } } - // databend::log::query - if let Err(error) = - InterpreterQueryLog::log_finish(ctx, now, error, has_profiles, &query_profiles) + if query_log_is_enabled + && let Err(error) = + InterpreterQueryLog::log_finish(ctx, now, error, has_profiles, &query_profiles) { error!("Failed to log query finish: {:?}", error) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use concurrent_queue::ConcurrentQueue; + use databend_common_base::runtime::CaptureLogSettings; + use databend_common_base::runtime::ThreadTracker; + use databend_common_tracing::HistoryConfig; + use databend_common_tracing::HistoryTableConfig; + use log::LevelFilter; + + use super::captured_info_log_enabled; + use super::history_table_log_enabled; + + #[test] + fn test_history_table_log_enabled() { + let mut history = HistoryConfig::default(); + history.tables.push(HistoryTableConfig { + table_name: "QUERY_HISTORY".to_string(), + ..Default::default() + }); + + assert!(!history_table_log_enabled(&history, "query_history")); + + history.on = true; + assert!(history_table_log_enabled(&history, "query_history")); + assert!(!history_table_log_enabled(&history, "profile_history")); + } + + #[test] + fn test_captured_info_log_enabled() { + assert!(!captured_info_log_enabled()); + + { + let mut payload = ThreadTracker::new_tracking_payload(); + payload.capture_log_settings = Some(CaptureLogSettings::capture_off()); + let _guard = ThreadTracker::tracking(payload); + assert!(!captured_info_log_enabled()); + } + + { + let mut payload = ThreadTracker::new_tracking_payload(); + payload.capture_log_settings = Some(CaptureLogSettings::capture_query( + LevelFilter::Info, + Arc::new(ConcurrentQueue::unbounded()), + )); + let _guard = ThreadTracker::tracking(payload); + assert!(captured_info_log_enabled()); + } + } +} diff --git a/src/query/service/src/interpreters/common/query_log.rs b/src/query/service/src/interpreters/common/query_log.rs index fe69aa0859819..0d56896528843 100644 --- a/src/query/service/src/interpreters/common/query_log.rs +++ b/src/query/service/src/interpreters/common/query_log.rs @@ -34,6 +34,7 @@ use log::info; use serde_json; use serde_json::Value; +use super::query_log_enabled; use crate::sessions::QueryContext; use crate::sessions::TableContextAuthorization; use crate::sessions::TableContextCluster; @@ -123,25 +124,31 @@ fn resource_usage_query_log(stats: IoStatsSnapshot, profiles: &[PlanProfile]) -> } impl InterpreterQueryLog { + pub fn enabled() -> bool { + query_log_enabled() + } + fn write_log(mut event: QueryLogElement) -> Result<()> { - // log the query event in the system_history.query_history table + // Normal callers are gated by `enabled()`. The log facade's target check is + // not specific enough to tell whether this query sink is configured. let event_str = serde_json::to_string(&event)?; info!(target: "databend::log::query", "{}", event_str); - // log the query event in `query-details` log file - // remove some fields to keep tidy in the log file + // Remove verbose fields from the query-details file. event.session_settings.clear(); event.sql_user_quota.clear(); event.sql_user_privileges.clear(); let event_str = serde_json::to_string(&event)?; info!(target: "databend::log::query::file", "{}", event_str); - // log the query event in the system log info!("query: {} becomes {:?}", event.query_id, event.log_type); Ok(()) } pub fn fail_to_start(ctx: Arc, err: ErrorCode) { + if !Self::enabled() { + return; + } InterpreterQueryLog::log_start(&ctx, SystemTime::now(), Some(err)) .unwrap_or_else(|e| error!("fail to write query_log {:?}", e)); } diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index 427784b0e7c3d..9a7138e433894 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -52,6 +52,7 @@ use crate::interpreters::ShowPublicKeysInterpreter; use crate::interpreters::UnsetObjectTagsInterpreter; use crate::interpreters::access::Accessor; use crate::interpreters::access_log::AccessLogger; +use crate::interpreters::common::access_log_enabled; use crate::interpreters::interpreter_add_warehouse_cluster::AddWarehouseClusterInterpreter; use crate::interpreters::interpreter_assign_warehouse_nodes::AssignWarehouseNodesInterpreter; use crate::interpreters::interpreter_catalog_drop::DropCatalogInterpreter; @@ -153,9 +154,11 @@ impl InterpreterFactory { error!("Access.denied(v2): {:?}", e); } })?; - let mut access_logger = AccessLogger::create(ctx.clone()); - access_logger.log(plan); - access_logger.output(); + if access_log_enabled() { + let mut access_logger = AccessLogger::create(ctx.clone()); + access_logger.log(plan); + access_logger.output(); + } Self::get_warehouses_interpreter(ctx, plan, Self::get_inner) } diff --git a/src/query/service/src/pipelines/executor/query_pipeline_executor.rs b/src/query/service/src/pipelines/executor/query_pipeline_executor.rs index f950395b379bf..c22e0508444fd 100644 --- a/src/query/service/src/pipelines/executor/query_pipeline_executor.rs +++ b/src/query/service/src/pipelines/executor/query_pipeline_executor.rs @@ -286,40 +286,60 @@ impl QueryPipelineExecutor { self.start_executor_daemon()?; - let mut thread_join_handles = self.execute_threads(self.threads_num); + // Pulling and complete pipelines call this from a dedicated OS thread, which can also act + // as one worker. Generic callers may invoke the executor from a Tokio runtime, where a + // processor using Tokio's blocking APIs would panic if it ran inline. + let execute_inline = tokio::runtime::Handle::try_current().is_err(); + let spawned_threads = self.threads_num - usize::from(execute_inline); + let mut thread_join_handles = self.execute_threads(spawned_threads); + + if execute_inline { + let inline_thread_num = spawned_threads; + let inline_thread_name = format!("PipelineExecutor-{}", inline_thread_num); + let inline_span = + Span::enter_with_local_parent("QueryPipelineExecutor::execute_threads") + .with_property(|| ("thread_name", inline_thread_name)); + let thread_res = self.execute_thread(inline_thread_num, inline_span); + self.check_thread_result(thread_res)?; + } while let Some(join_handle) = thread_join_handles.pop() { let thread_res = join_handle.join().flatten(); + self.check_thread_result(thread_res)?; + } - { - let finished_error_guard = self.finished_error.lock(); - if let Some(error) = finished_error_guard.as_ref() { - let may_error = error.clone(); - drop(finished_error_guard); - - let profiling = self.fetch_plans_profile(true); - self.on_finished(ExecutionInfo::create(Err(may_error.clone()), profiling))?; - return Err(may_error); - } - } + if let Err(error) = self.graph.assert_finished_graph() { + let profiling = self.fetch_plans_profile(true); + self.on_finished(ExecutionInfo::create(Err(error.clone()), profiling))?; + return Err(error); + } + + let profiling = self.fetch_plans_profile(true); + self.on_finished(ExecutionInfo::create(Ok(()), profiling))?; + Ok(()) + } + + fn check_thread_result(&self, thread_res: Result<()>) -> Result<()> { + { + let finished_error_guard = self.finished_error.lock(); + if let Some(error) = finished_error_guard.as_ref() { + let may_error = error.clone(); + drop(finished_error_guard); - // We will ignore the abort query error, because returned by finished_error if abort query. - if matches!(&thread_res, Err(error) if error.code() != ErrorCode::ABORTED_QUERY) { - let may_error = thread_res.unwrap_err(); let profiling = self.fetch_plans_profile(true); self.on_finished(ExecutionInfo::create(Err(may_error.clone()), profiling))?; return Err(may_error); } } - if let Err(error) = self.graph.assert_finished_graph() { + // We will ignore the abort query error, because returned by finished_error if abort query. + if matches!(&thread_res, Err(error) if error.code() != ErrorCode::ABORTED_QUERY) { + let may_error = thread_res.unwrap_err(); let profiling = self.fetch_plans_profile(true); - self.on_finished(ExecutionInfo::create(Err(error.clone()), profiling))?; - return Err(error); + self.on_finished(ExecutionInfo::create(Err(may_error.clone()), profiling))?; + return Err(may_error); } - let profiling = self.fetch_plans_profile(true); - self.on_finished(ExecutionInfo::create(Ok(()), profiling))?; Ok(()) } @@ -417,28 +437,32 @@ impl QueryPipelineExecutor { let span = Span::enter_with_local_parent("QueryPipelineExecutor::execute_threads") .with_property(|| ("thread_name", name.clone())); - thread_join_handles.push(Thread::named_spawn(Some(name), move || unsafe { - let _g = span.set_local_parent(); - let this_clone = this.clone(); - let try_result = catch_unwind(|| this_clone.execute_single_thread(thread_num)); - - // finish the pipeline executor when has error or panic - if let Err(cause) = try_result.flatten() { - span.with_property(|| ("error", "true")).add_properties(|| { - [ - ("error.type", cause.code().to_string()), - ("error.message", cause.display_text()), - ] - }); - this.finish(Some(cause)); - } - - Ok(()) + thread_join_handles.push(Thread::named_spawn(Some(name), move || { + this.execute_thread(thread_num, span) })); } thread_join_handles } + fn execute_thread(self: &Arc, thread_num: usize, span: Span) -> Result<()> { + let _g = span.set_local_parent(); + let this = self.clone(); + let try_result = catch_unwind(|| unsafe { this.execute_single_thread(thread_num) }); + + // Finish the pipeline executor when a worker returns an error or panics. + if let Err(cause) = try_result.flatten() { + span.with_property(|| ("error", "true")).add_properties(|| { + [ + ("error.type", cause.code().to_string()), + ("error.message", cause.display_text()), + ] + }); + self.finish(Some(cause)); + } + + Ok(()) + } + /// # Safety /// /// Method is thread unsafe and require thread safe call diff --git a/src/query/service/src/servers/mysql/mysql_interactive_worker.rs b/src/query/service/src/servers/mysql/mysql_interactive_worker.rs index 183ac64fa0741..3cda835d755e5 100644 --- a/src/query/service/src/servers/mysql/mysql_interactive_worker.rs +++ b/src/query/service/src/servers/mysql/mysql_interactive_worker.rs @@ -23,6 +23,7 @@ use databend_common_base::base::mask_connection_info; use databend_common_base::runtime::MemStat; use databend_common_base::runtime::ThreadTracker; use databend_common_base::runtime::TrackingPayloadExt; +use databend_common_base::runtime::spawn; use databend_common_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; @@ -474,7 +475,9 @@ impl InteractiveWorkerBase { )> { let instant = Instant::now(); - let query_result = context.try_spawn({ + // The connection already runs on the shared MySQL query executor. Keep the task boundary + // without creating and destroying a dedicated Tokio runtime for every statement. + let query_result = spawn({ let ctx = context.clone(); async move { let mut data_stream = interpreter.execute(ctx.clone()).await?; @@ -491,11 +494,11 @@ impl InteractiveWorkerBase { Ok::<_, ErrorCode>(intercepted_stream.boxed()) } .in_span(Span::enter_with_local_parent(func_path!())) - })?; + }); let query_result = query_result.await.map_err_to_code( ErrorCode::TokioError, - || "Cannot join handle from context's runtime", + || "Cannot join handle from MySQL query executor runtime", )?; let reporter = Box::new(ContextProgressReporter::new(context.clone(), instant)) as Box; diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs index 105ed602618e7..33fb76bdd3b41 100644 --- a/src/query/service/src/sessions/query_ctx.rs +++ b/src/query/service/src/sessions/query_ctx.rs @@ -26,6 +26,7 @@ use std::collections::HashSet; use std::collections::VecDeque; use std::future::Future; use std::sync::Arc; +use std::sync::Once; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; @@ -192,6 +193,7 @@ pub struct QueryContext { partition_queue: Arc>>, shared: Arc, query_settings: Arc, + session_settings_initialized: Once, fragment_id: FragmentId, // Used by synchronized generate aggregating indexes when new data written. written_segment_locs: SegmentLocationsState, @@ -219,6 +221,7 @@ impl QueryContext { mysql_version: format!("{MYSQL_VERSION}-{}", shared.version.commit_detail), shared, query_settings, + session_settings_initialized: Once::new(), fragment_id: Default::default(), written_segment_locs: Default::default(), read_block_thresholds: Default::default(), diff --git a/src/query/service/src/sessions/query_ctx/context.rs b/src/query/service/src/sessions/query_ctx/context.rs index c39dc1e2a1dc4..bd29eacadfac3 100644 --- a/src/query/service/src/sessions/query_ctx/context.rs +++ b/src/query/service/src/sessions/query_ctx/context.rs @@ -323,8 +323,8 @@ impl TableContextSettings for QueryContext { } fn get_settings(&self) -> Arc { - if self.shared.query_settings.is_changed() - && self.shared.query_settings.query_level_change() + if self.shared.query_settings.query_level_change() + && self.shared.query_settings.is_changed() { let shared_settings = self.shared.query_settings.changes(); if self.get_session_settings().is_changed() { @@ -342,10 +342,10 @@ impl TableContextSettings for QueryContext { } } } else { - unsafe { + self.session_settings_initialized.call_once(|| unsafe { self.query_settings - .unchecked_apply_changes(self.get_session_settings().changes()) - } + .unchecked_apply_changes(self.get_session_settings().changes()); + }); } self.query_settings.clone() diff --git a/src/query/service/src/sessions/session_ctx.rs b/src/query/service/src/sessions/session_ctx.rs index c396ecb5bb507..af510e9a7b992 100644 --- a/src/query/service/src/sessions/session_ctx.rs +++ b/src/query/service/src/sessions/session_ctx.rs @@ -37,6 +37,48 @@ use parking_lot::RwLock; use crate::sessions::QueryContextShared; +#[derive(Default)] +struct QueryIdsResults { + entries: Vec<(String, Option)>, + indices: HashMap, +} + +impl QueryIdsResults { + fn get(&self, query_id: &str) -> Option { + let index = self.indices.get(&query_id.to_ascii_lowercase())?; + self.entries.get(*index)?.1.clone() + } + + fn update(&mut self, query_id: String, value: Option) { + let normalized_id = query_id.to_ascii_lowercase(); + if let Some(index) = self.indices.get(&normalized_id).copied() { + if let Some(value) = value { + self.entries[index] = (query_id, Some(value)); + } + return; + } + + let index = self.entries.len(); + self.entries.push((query_id, value)); + self.indices.insert(normalized_id, index); + } + + fn last(&self, index: i32) -> Option { + let len = self.entries.len(); + let index = if index < 0 { len as i32 + index } else { index }; + + if len < 1 || index < 0 || index > (len - 1) as i32 { + return None; + } + + Some(self.entries[index as usize].0.clone()) + } + + fn ids(&self) -> HashSet { + HashSet::from_iter(self.entries.iter().map(|result| result.0.clone())) + } +} + pub struct SessionContext { abort: AtomicBool, settings: Arc, @@ -71,7 +113,7 @@ pub struct SessionContext { query_context_shared: RwLock>, /// We store `query_id -> query_result_cache_key` to session context, so that we can fetch /// query result through previous query_id easily. - query_ids_results: RwLock)>>, + query_ids_results: RwLock, // Used in set variables inside session variables: Arc>>, typ: SessionType, @@ -306,49 +348,19 @@ impl SessionContext { } pub fn get_query_result_cache_key(&self, query_id: &str) -> Option { - let lock = self.query_ids_results.read(); - for (qid, result_cache_key) in (*lock).iter().rev() { - if qid.eq_ignore_ascii_case(query_id) { - return result_cache_key.clone(); - } - } - None + self.query_ids_results.read().get(query_id) } pub fn update_query_ids_results(&self, query_id: String, value: Option) { - let mut lock = self.query_ids_results.write(); - // Here we use reverse iteration, as it is not common to modify elements from earlier. - for (idx, (qid, _)) in (*lock).iter().rev().enumerate() { - if qid.eq_ignore_ascii_case(&query_id) { - // update value iff value is some. - if let Some(v) = value { - (*lock)[idx] = (query_id, Some(v)) - } - return; - } - } - lock.push((query_id, value)) + self.query_ids_results.write().update(query_id, value) } pub fn get_last_query_id(&self, index: i32) -> Option { - let lock = self.query_ids_results.read(); - let query_ids_len = lock.len(); - let idx = if index < 0 { - query_ids_len as i32 + index - } else { - index - }; - - if query_ids_len < 1 || idx < 0 || idx > (query_ids_len - 1) as i32 { - return None; - } - - Some((*lock)[idx as usize].0.clone()) + self.query_ids_results.read().last(index) } pub fn get_query_id_history(&self) -> HashSet { - let lock = self.query_ids_results.read(); - HashSet::from_iter(lock.iter().map(|result| result.clone().0)) + self.query_ids_results.read().ids() } pub fn txn_mgr(&self) -> TxnManagerRef { @@ -414,3 +426,31 @@ impl SessionContext { *self.current_workload_group.write() = Some(workload_group) } } + +#[cfg(test)] +mod tests { + use super::QueryIdsResults; + + #[test] + fn test_query_ids_results_updates_the_matching_entry() { + let mut history = QueryIdsResults::default(); + history.update("first".to_string(), None); + history.update("second".to_string(), None); + history.update("SECOND".to_string(), Some("cache-key".to_string())); + + assert_eq!(history.entries.len(), 2); + assert_eq!(history.get("second"), Some("cache-key".to_string())); + assert_eq!(history.last(-1), Some("SECOND".to_string())); + assert_eq!(history.last(0), Some("first".to_string())); + } + + #[test] + fn test_query_ids_results_ignores_none_for_an_existing_entry() { + let mut history = QueryIdsResults::default(); + history.update("query".to_string(), Some("existing-cache-key".to_string())); + history.update("QUERY".to_string(), None); + + assert_eq!(history.entries.len(), 1); + assert_eq!(history.get("query"), Some("existing-cache-key".to_string())); + } +} diff --git a/src/query/service/tests/it/pipelines/executor/pipeline_executor.rs b/src/query/service/tests/it/pipelines/executor/pipeline_executor.rs index 764dba8e94479..197d4842cf6e6 100644 --- a/src/query/service/tests/it/pipelines/executor/pipeline_executor.rs +++ b/src/query/service/tests/it/pipelines/executor/pipeline_executor.rs @@ -14,8 +14,10 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use databend_common_base::runtime::Thread; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_expression::DataBlock; @@ -95,6 +97,59 @@ async fn test_always_call_on_finished() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_execute_inside_and_outside_tokio_runtime() -> anyhow::Result<()> { + let fixture = TestFixture::setup().await?; + let ctx = fixture.new_query_ctx().await?; + let settings = ExecutorSettings { + query_id: Arc::new("worker-mode-test".to_string()), + max_execute_time_in_seconds: Default::default(), + enable_queries_executor: false, + max_threads: 1, + executor_node_id: "".to_string(), + perf_event_groups: vec![], + }; + + let (tokio_finished, tokio_executor) = + create_closed_source_executor(ctx.clone(), settings.clone())?; + tokio_executor.execute()?; + assert_eq!(tokio_finished.load(Ordering::SeqCst), 1); + + let (inline_finished, inline_executor) = create_closed_source_executor(ctx, settings)?; + Thread::spawn(move || inline_executor.execute()) + .join() + .expect("inline executor thread must not panic")?; + assert_eq!(inline_finished.load(Ordering::SeqCst), 1); + + Ok(()) +} + +fn create_closed_source_executor( + ctx: Arc, + settings: ExecutorSettings, +) -> Result<(Arc, Arc)> { + let mut pipeline = Pipeline::create(); + let (source_senders, source_pipe) = create_source_pipe(ctx, 1)?; + let (_sink_receivers, sink_pipe) = create_sink_pipe(1)?; + drop(source_senders); + + pipeline.add_pipe(source_pipe); + pipeline.add_pipe(sink_pipe); + pipeline.set_max_threads(1); + + let finished = Arc::new(AtomicUsize::new(0)); + pipeline.set_on_finished({ + let finished = finished.clone(); + move |_info: &ExecutionInfo| { + finished.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + }); + + let executor = QueryPipelineExecutor::create(pipeline, settings)?; + Ok((finished, executor)) +} + fn create_pipeline() -> (Arc, Pipeline) { let called_finished = Arc::new(AtomicBool::new(false)); let mut pipeline = Pipeline::create(); diff --git a/src/query/service/tests/it/sessions/session_setting.rs b/src/query/service/tests/it/sessions/session_setting.rs index 9b4579f3f7107..5185ed79fa407 100644 --- a/src/query/service/tests/it/sessions/session_setting.rs +++ b/src/query/service/tests/it/sessions/session_setting.rs @@ -13,6 +13,7 @@ // limitations under the License. use databend_common_catalog::session_type::SessionType; +use databend_query::sessions::TableContextSettings; use databend_query::test_kits::ConfigBuilder; use databend_query::test_kits::TestFixture; @@ -40,6 +41,29 @@ async fn test_session_setting() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_query_context_snapshots_session_settings() -> anyhow::Result<()> { + let fixture = TestFixture::setup().await?; + let session = fixture.default_session(); + session + .get_settings() + .set_setting("max_threads".to_string(), "2".to_string())?; + + let first_ctx = fixture.new_query_ctx().await?; + assert_eq!(first_ctx.get_settings().get_max_threads()?, 2); + + session + .get_settings() + .set_setting("max_threads".to_string(), "3".to_string())?; + + assert_eq!(first_ctx.get_settings().get_max_threads()?, 2); + + let next_ctx = fixture.new_query_ctx().await?; + assert_eq!(next_ctx.get_settings().get_max_threads()?, 3); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_session_setting_override() -> anyhow::Result<()> { // Setup. diff --git a/src/query/service/tests/it/storages/fuse/operations/prewhere.rs b/src/query/service/tests/it/storages/fuse/operations/prewhere.rs index 0c089ee730095..1190103e23997 100644 --- a/src/query/service/tests/it/storages/fuse/operations/prewhere.rs +++ b/src/query/service/tests/it/storages/fuse/operations/prewhere.rs @@ -329,6 +329,7 @@ async fn prepare_prewhere_data() -> Result { let part = FuseBlockPartInfo { location: "test_block".to_string(), bloom_filter_index_location: None, + file_size: parquet_bytes.len() as u64, bloom_filter_index_size: 0, create_on: None, nums_rows: num_rows, diff --git a/src/query/storages/common/cache/Cargo.toml b/src/query/storages/common/cache/Cargo.toml index 07799905a83a5..2852d0145edcf 100644 --- a/src/query/storages/common/cache/Cargo.toml +++ b/src/query/storages/common/cache/Cargo.toml @@ -14,6 +14,7 @@ databend-common-cache = { workspace = true } databend-common-catalog = { workspace = true } databend-common-config = { workspace = true } databend-common-exception = { workspace = true } +databend-common-expression = { workspace = true } databend-common-metrics = { workspace = true } databend-storages-common-index = { workspace = true } databend-storages-common-table-meta = { workspace = true } diff --git a/src/query/storages/common/cache/src/caches.rs b/src/query/storages/common/cache/src/caches.rs index 39bf1da621be3..24885216d4090 100644 --- a/src/query/storages/common/cache/src/caches.rs +++ b/src/query/storages/common/cache/src/caches.rs @@ -18,6 +18,7 @@ use std::time::Instant; use arrow::array::ArrayRef; use databend_common_cache::MemSized; +use databend_common_expression::DataBlock; use crate::CacheAccessor; use crate::InMemoryLruCache; @@ -69,10 +70,53 @@ pub type PrunePartitionsCache = InMemoryLruCache<(PartStatistics, Partitions)>; pub type IcebergTableCache = InMemoryLruCache<(Arc, AtomicBool, Instant)>; -/// In memory object cache of table column array -pub type ColumnArrayCache = InMemoryLruCache; -pub type ArrayRawDataUncompressedSize = usize; -pub type SizedColumnArray = (ArrayRef, ArrayRawDataUncompressedSize); +/// In-memory cache of decoded table data. +pub type TableDataCache = InMemoryLruCache; + +pub enum TableDataCacheValue { + ColumnArray(ArrayRef), + DataBlock(DataBlock), +} + +pub struct TableDataCacheEntry { + value: TableDataCacheValue, + memory_size: usize, +} + +impl TableDataCacheEntry { + pub fn from_column_array(value: ArrayRef, memory_size: usize) -> Self { + Self { + value: TableDataCacheValue::ColumnArray(value), + memory_size, + } + } + + pub fn from_data_block(value: DataBlock) -> Self { + let memory_size = value.memory_size(); + Self { + value: TableDataCacheValue::DataBlock(value), + memory_size, + } + } + + pub fn as_column_array(&self) -> Option<&ArrayRef> { + match &self.value { + TableDataCacheValue::ColumnArray(value) => Some(value), + TableDataCacheValue::DataBlock(_) => None, + } + } + + pub fn as_data_block(&self) -> Option<&DataBlock> { + match &self.value { + TableDataCacheValue::ColumnArray(_) => None, + TableDataCacheValue::DataBlock(value) => Some(value), + } + } + + pub fn memory_size(&self) -> usize { + self.memory_size + } +} // Bind Type of cached objects to Caches // @@ -359,10 +403,10 @@ impl From<(PartStatistics, Partitions)> for CacheValue<(PartStatistics, Partitio } } -impl From for CacheValue { - fn from(value: SizedColumnArray) -> Self { +impl From for CacheValue { + fn from(value: TableDataCacheEntry) -> Self { CacheValue { - mem_bytes: value.1, + mem_bytes: value.memory_size, inner: Arc::new(value), } } @@ -384,3 +428,35 @@ impl MemSized for CacheValue { self.mem_bytes } } + +#[cfg(test)] +mod tests { + use arrow::array::Int32Array; + + use super::*; + use crate::CacheAccessor; + + #[test] + fn test_table_data_cache_mixed_entries() { + let cache = TableDataCache::with_bytes_capacity("test_table_data".to_string(), 1024); + + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + cache.insert( + "column".to_string(), + TableDataCacheEntry::from_column_array(array.clone(), 128), + ); + let cached_array = cache.get("column").unwrap(); + assert!(cached_array.as_data_block().is_none()); + assert_eq!(cached_array.as_column_array().unwrap().len(), array.len()); + assert_eq!(cached_array.memory_size(), 128); + + cache.insert( + "block".to_string(), + TableDataCacheEntry::from_data_block(DataBlock::empty_with_rows(3)), + ); + let cached_block = cache.get("block").unwrap(); + assert!(cached_block.as_column_array().is_none()); + assert_eq!(cached_block.as_data_block().unwrap().num_rows(), 3); + assert_eq!(cached_block.memory_size(), 0); + } +} diff --git a/src/query/storages/common/cache/src/manager.rs b/src/query/storages/common/cache/src/manager.rs index 39914f449644c..45ec3bfcecbfd 100644 --- a/src/query/storages/common/cache/src/manager.rs +++ b/src/query/storages/common/cache/src/manager.rs @@ -36,7 +36,6 @@ use crate::caches::BlockMetaCache; use crate::caches::BloomIndexFilterCache; use crate::caches::BloomIndexMetaCache; use crate::caches::CacheValue; -use crate::caches::ColumnArrayCache; use crate::caches::ColumnDataCache; use crate::caches::ColumnOrientedSegmentInfoCache; use crate::caches::CompactSegmentInfoCache; @@ -48,6 +47,7 @@ use crate::caches::PrunePartitionsCache; use crate::caches::SegmentBlockMetasCache; use crate::caches::SpatialIndexFileCache; use crate::caches::SpatialIndexMetaCache; +use crate::caches::TableDataCache; use crate::caches::TableSnapshotCache; use crate::caches::TableSnapshotStatisticCache; use crate::caches::VectorIndexFileCache; @@ -118,7 +118,7 @@ pub struct CacheManager { virtual_column_meta_cache: CacheSlot, prune_partitions_cache: CacheSlot, parquet_meta_data_cache: CacheSlot, - in_memory_table_data_cache: CacheSlot, + in_memory_table_data_cache: CacheSlot, segment_block_metas_cache: CacheSlot, block_meta_cache: CacheSlot, @@ -835,7 +835,7 @@ impl CacheManager { self.get_hybrid_cache(self.column_data_cache.get()) } - pub fn get_table_data_array_cache(&self) -> Option { + pub fn get_table_data_cache(&self) -> Option { self.in_memory_table_data_cache.get() } @@ -1023,6 +1023,8 @@ mod tests { use super::*; use crate::ColumnData; + use crate::TableDataCacheEntry; + fn config_with_disk_cache_enabled(cache_path: &str) -> CacheConfig { CacheConfig { data_cache_storage: CacheStorageTypeInnerConfig::Disk, @@ -1259,9 +1261,12 @@ mod tests { // ----- POPULATE BASIC CACHES ----- // 1. Populate in-memory table data cache - let in_memory_table_data_cache = cache_manager.get_table_data_array_cache().unwrap(); + let in_memory_table_data_cache = cache_manager.get_table_data_cache().unwrap(); let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); - in_memory_table_data_cache.insert("not matter".to_string(), (array, 1)); + in_memory_table_data_cache.insert( + "not matter".to_string(), + TableDataCacheEntry::from_column_array(array, 1), + ); // 2. Populate segment block metas cache let segment_block_metas_cache = cache_manager.get_segment_block_metas_cache().unwrap(); @@ -1341,7 +1346,7 @@ mod tests { // ----- VERIFY INITIAL CACHE STATE ----- // Verify all caches are correctly populated - assert!(!cache_manager.get_table_data_array_cache().is_empty()); + assert!(!cache_manager.get_table_data_cache().is_empty()); assert!( !cache_manager .get_segment_block_metas_cache() diff --git a/src/query/storages/fuse/Cargo.toml b/src/query/storages/fuse/Cargo.toml index 0c6d44822b9f0..20f843a059f3f 100644 --- a/src/query/storages/fuse/Cargo.toml +++ b/src/query/storages/fuse/Cargo.toml @@ -25,6 +25,7 @@ databend-common-pipeline-transforms = { workspace = true } databend-common-sql = { workspace = true } databend-common-statistics = { workspace = true } databend-common-storage = { workspace = true } +databend-common-storages-parquet = { workspace = true } databend-common-tracing = { workspace = true } databend-common-users = { workspace = true } databend-enterprise-fail-safe = { workspace = true } diff --git a/src/query/storages/fuse/src/fuse_part.rs b/src/query/storages/fuse/src/fuse_part.rs index e18e306cb6ef0..6dcba37377073 100644 --- a/src/query/storages/fuse/src/fuse_part.rs +++ b/src/query/storages/fuse/src/fuse_part.rs @@ -45,6 +45,8 @@ pub struct FuseBlockPartInfo { pub create_on: Option>, pub nums_rows: usize, + #[serde(default)] + pub file_size: u64, pub columns_meta: HashMap, pub columns_stat: Option>, pub compression: Compression, @@ -96,6 +98,7 @@ impl FuseBlockPartInfo { bloom_filter_index_size, create_on, columns_meta, + file_size: 0, nums_rows: rows_count as usize, compression, sort_min_max, diff --git a/src/query/storages/fuse/src/io/read/block/block_reader.rs b/src/query/storages/fuse/src/io/read/block/block_reader.rs index 01dc040fab256..0a702bbccfbdc 100644 --- a/src/query/storages/fuse/src/io/read/block/block_reader.rs +++ b/src/query/storages/fuse/src/io/read/block/block_reader.rs @@ -272,7 +272,7 @@ impl BlockReadContext { let read_from_in_mem_cache_array: usize = block_read_res .cached_column_array .iter() - .map(|(_, sized_array)| sized_array.1) + .map(|(_, cache_entry)| cache_entry.memory_size()) .sum(); cache_metrics.add_cache_metrics( diff --git a/src/query/storages/fuse/src/io/read/block/block_reader_merge_io.rs b/src/query/storages/fuse/src/io/read/block/block_reader_merge_io.rs index cdb61afcc436a..d8b69991e5edc 100644 --- a/src/query/storages/fuse/src/io/read/block/block_reader_merge_io.rs +++ b/src/query/storages/fuse/src/io/read/block/block_reader_merge_io.rs @@ -19,18 +19,18 @@ use bytes::Bytes; use databend_common_exception::Result; use databend_common_expression::ColumnId; use databend_storages_common_cache::ColumnData; -use databend_storages_common_cache::SizedColumnArray; +use databend_storages_common_cache::TableDataCacheEntry; use databend_storages_common_io::MergeIOReadResult; use enum_as_inner::EnumAsInner; use opendal::Buffer; type CachedColumnData = Vec<(ColumnId, Arc)>; -type CachedColumnArray = Vec<(ColumnId, Arc)>; +type CachedColumnArray = Vec<(ColumnId, Arc)>; #[derive(EnumAsInner, Clone)] pub enum DataItem<'a> { RawData(Buffer), - ColumnArray(&'a Arc), + ColumnArray(&'a Arc), } pub struct BlockReadResult { diff --git a/src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs b/src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs index 1d0be8fa3cb0c..b67ec9aae9016 100644 --- a/src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs +++ b/src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs @@ -64,7 +64,7 @@ impl BlockReadContext { let mut ranges = vec![]; // for async read, try using table data cache (if enabled in settings) let column_data_cache = CacheManager::instance().get_column_data_cache(); - let column_array_cache = CacheManager::instance().get_table_data_array_cache(); + let table_data_cache = CacheManager::instance().get_table_data_cache(); let mut cached_column_data = vec![]; let mut cached_column_array = vec![]; @@ -84,14 +84,16 @@ impl BlockReadContext { // first, check in memory table data cache // column_array_cache - if let Some(cache_array) = column_array_cache.get_sized(&column_cache_key, len) { - // Record bytes scanned from memory cache (table data only) - Profile::record_usize_profile( - ProfileStatisticsName::ScanBytesFromMemory, - len as usize, - ); - cached_column_array.push((*column_id, cache_array)); - continue; + if let Some(cache_entry) = table_data_cache.get_sized(&column_cache_key, len) { + if cache_entry.as_column_array().is_some() { + // Record bytes scanned from memory cache (table data only) + Profile::record_usize_profile( + ProfileStatisticsName::ScanBytesFromMemory, + len as usize, + ); + cached_column_array.push((*column_id, cache_entry)); + continue; + } } // and then, check on disk table data cache diff --git a/src/query/storages/fuse/src/io/read/block/parquet/mod.rs b/src/query/storages/fuse/src/io/read/block/parquet/mod.rs index aaeb22d5291b3..4d69172714515 100644 --- a/src/query/storages/fuse/src/io/read/block/parquet/mod.rs +++ b/src/query/storages/fuse/src/io/read/block/parquet/mod.rs @@ -13,11 +13,14 @@ // limitations under the License. use std::collections::HashMap; +use std::fmt::Write; use arrow_array::Array; use arrow_array::ArrayRef; use arrow_array::RecordBatch; use arrow_array::StructArray; +use databend_common_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; use databend_common_catalog::plan::Projection; use databend_common_exception::ErrorCode; use databend_common_expression::BlockEntry; @@ -27,9 +30,11 @@ use databend_common_expression::FilterVisitor; use databend_common_expression::TableDataType; use databend_common_expression::TableSchema; use databend_common_expression::Value; +use databend_common_expression::types::DataType; use databend_common_expression::visitor::ValueVisitor; use databend_storages_common_cache::CacheAccessor; use databend_storages_common_cache::CacheManager; +use databend_storages_common_cache::TableDataCacheEntry; use databend_storages_common_cache::TableDataCacheKey; use databend_storages_common_table_meta::meta::ColumnMeta; use databend_storages_common_table_meta::meta::Compression; @@ -62,6 +67,124 @@ impl BlockReader { ) } + pub fn page_range_bitmap( + part: &FuseBlockPartInfo, + ) -> Option { + part.range().map(|range| { + RowSelection::from_range( + part.nums_rows, + range.start.saturating_mul(part.page_size()), + range.end.saturating_mul(part.page_size()), + ) + .bitmap + }) + } + + pub fn page_range_data_cache_key(&self, part: &FuseBlockPartInfo) -> Option { + let range = part.range()?; + let root = self.operator.info().root(); + let mut key = String::with_capacity(root.len() + part.location.len() + 256); + write!( + key, + "fuse-page-data-v1|{}:{}|{}:{}|{}|{}|{}|{}:{}", + root.len(), + root, + part.location.len(), + part.location, + part.file_size, + part.nums_rows, + part.page_size(), + range.start, + range.end, + ) + .ok()?; + + for (field, column_node) in self + .projected_schema + .fields + .iter() + .zip(self.project_column_nodes.iter()) + { + if column_node.is_nested || column_node.leaf_column_ids.as_slice() != [field.column_id] + { + return None; + } + let column_meta = part.columns_meta.get(&field.column_id)?; + let (offset, len) = column_meta.offset_length(); + let data_type: DataType = field.data_type().into(); + write!( + key, + "|{}:{}:{}:{:?}", + field.column_id, offset, len, data_type + ) + .ok()?; + } + Some(key) + } + + pub fn cached_page_range_data(&self, key: &str) -> Option { + let cache = CacheManager::instance().get_table_data_cache(); + let cache_entry = cache.get(key)?; + let data_block = cache_entry.as_data_block()?; + let memory_size = cache_entry.memory_size(); + Profile::record_usize_profile(ProfileStatisticsName::ScanBytesFromMemory, memory_size); + self.ctx + .get_data_cache_metrics() + .add_cache_metrics(0, 0, memory_size); + Some(data_block.clone()) + } + + pub fn cache_page_range_data(&self, key: String, data_block: DataBlock) -> DataBlock { + if !self.put_cache { + return data_block; + } + let Some(cache) = CacheManager::instance().get_table_data_cache() else { + return data_block; + }; + let cache_entry = cache.insert(key, TableDataCacheEntry::from_data_block(data_block)); + cache_entry.as_data_block().unwrap().clone() + } + + pub fn deserialize_parquet_record_batch( + &self, + part: &FuseBlockPartInfo, + record_batch: &RecordBatch, + ) -> databend_common_exception::Result { + let result_rows = record_batch.num_rows(); + if self.projected_schema.fields.is_empty() { + return Ok(DataBlock::empty_with_rows(result_rows)); + } + + if result_rows == 0 { + return Ok(DataBlock::empty_with_schema(&self.data_schema())); + } + + let mut entries = Vec::with_capacity(self.projected_schema.fields.len()); + let name_paths = column_name_paths(&self.projection, &self.original_schema); + for (((i, field), column_node), name_path) in self + .projected_schema + .fields + .iter() + .enumerate() + .zip(self.project_column_nodes.iter()) + .zip(name_paths.iter()) + { + let data_type = field.data_type().into(); + let exists = column_node + .leaf_column_ids + .iter() + .any(|column_id| part.columns_meta.contains_key(column_id)); + let value = if exists { + let arrow_array = column_by_name(record_batch, name_path); + Value::from_arrow_rs(arrow_array, &data_type)? + } else { + Value::Scalar(self.default_vals[i].clone()) + }; + entries.push(BlockEntry::new(value, || (data_type, result_rows))); + } + Ok(DataBlock::new(entries, result_rows)) + } + pub fn deserialize_parquet_chunks( &self, num_rows: usize, @@ -94,7 +217,7 @@ impl BlockReader { let name_paths = column_name_paths(&self.projection, &self.original_schema); let array_cache = if self.put_cache && !has_selection { - CacheManager::instance().get_table_data_array_cache() + CacheManager::instance().get_table_data_cache() } else { None }; @@ -133,7 +256,13 @@ impl BlockReader { let key = TableDataCacheKey::new(block_path, field.column_id, offset, len); let array_memory_size = arrow_array.get_array_memory_size(); - cache.insert(key.into(), (arrow_array.clone(), array_memory_size)); + cache.insert( + key.into(), + TableDataCacheEntry::from_column_array( + arrow_array.clone(), + array_memory_size, + ), + ); } } Value::from_arrow_rs(arrow_array, &data_type)? @@ -145,7 +274,10 @@ impl BlockReader { "unexpected nested field: nested leaf field hits cached", )); } - let mut value = Value::from_arrow_rs(cached.0.clone(), &data_type)?; + let cached_array = cached.as_column_array().ok_or_else(|| { + ErrorCode::StorageOther("unexpected data block entry in column array cache") + })?; + let mut value = Value::from_arrow_rs(cached_array.clone(), &data_type)?; if let Some(selection) = selection { let mut filter_visitor = FilterVisitor::new(&selection.bitmap); filter_visitor.visit_value(value)?; diff --git a/src/query/storages/fuse/src/io/read/block/parquet/row_selection.rs b/src/query/storages/fuse/src/io/read/block/parquet/row_selection.rs index 4981db3d9fd59..b58289ef13390 100644 --- a/src/query/storages/fuse/src/io/read/block/parquet/row_selection.rs +++ b/src/query/storages/fuse/src/io/read/block/parquet/row_selection.rs @@ -14,6 +14,7 @@ use databend_common_column::bitmap::utils::SlicesIterator; use databend_common_expression::types::Bitmap; +use databend_common_expression::types::MutableBitmap; use parquet::arrow::arrow_reader::RowSelector; /// A wrapper around parquet's `RowSelection` that also tracks the number of selected rows (bits set to 1 in the bitmap). @@ -38,6 +39,37 @@ impl RowSelection { bitmap, } } + + pub fn from_range(total_rows: usize, start: usize, end: usize) -> Self { + let start = start.min(total_rows); + let end = end.min(total_rows).max(start); + + let selected_rows = end - start; + let mut bitmap = MutableBitmap::with_capacity(total_rows); + bitmap.extend_constant(start, false); + bitmap.extend_constant(selected_rows, true); + bitmap.extend_constant(total_rows - end, false); + + let mut selectors = Vec::with_capacity(3); + if selected_rows == 0 { + if total_rows > 0 { + selectors.push(RowSelector::skip(total_rows)); + } + } else { + if start > 0 { + selectors.push(RowSelector::skip(start)); + } + selectors.push(RowSelector::select(selected_rows)); + if end < total_rows { + selectors.push(RowSelector::skip(total_rows - end)); + } + } + Self::new( + parquet::arrow::arrow_reader::RowSelection::from(selectors), + selected_rows, + bitmap.into(), + ) + } } impl From<&Bitmap> for RowSelection { @@ -136,4 +168,25 @@ mod tests { assert_eq!(selectors[0].row_count, 5); assert!(selectors[0].skip); } + + #[test] + fn test_row_selection_from_range() { + let row_selection = RowSelection::from_range(10, 2, 5); + let selectors: Vec<_> = row_selection.selection.iter().collect(); + + assert_eq!(row_selection.selected_rows, 3); + assert_eq!(row_selection.bitmap.iter().collect::>(), vec![ + false, false, true, true, true, false, false, false, false, false + ]); + assert_eq!(selectors.len(), 3); + assert_eq!((selectors[0].row_count, selectors[0].skip), (2, true)); + assert_eq!((selectors[1].row_count, selectors[1].skip), (3, false)); + assert_eq!((selectors[2].row_count, selectors[2].skip), (5, true)); + + let empty = RowSelection::from_range(10, 4, 4); + let selectors: Vec<_> = empty.selection.iter().collect(); + assert_eq!(empty.selected_rows, 0); + assert_eq!(selectors.len(), 1); + assert_eq!((selectors[0].row_count, selectors[0].skip), (10, true)); + } } diff --git a/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs b/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs index 98b4000a8c210..64020811423d6 100644 --- a/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs +++ b/src/query/storages/fuse/src/operations/read/fuse_rows_fetcher.rs @@ -73,6 +73,7 @@ pub fn row_fetch_processor( FuseStorageFormat::Parquet => { let read_settings = ReadSettings::from_ctx(&ctx)?; let max_threads = ctx.get_settings().get_max_threads()? as usize; + let source_parts = source.parts.partitions.clone(); let block_threshold = BlockThreshold { max_rows: ctx.get_settings().get_max_block_size()? as usize, max_bytes: ctx.get_settings().get_max_block_bytes()? as usize, @@ -92,6 +93,7 @@ pub fn row_fetch_processor( read_settings, max_threads, io_semaphore.clone(), + &source_parts, ), need_wrap_nullable, fetched_data_types.clone(), diff --git a/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs b/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs index 375acb536a130..bfdc39999a09f 100644 --- a/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs +++ b/src/query/storages/fuse/src/operations/read/parquet_rows_fetcher.rs @@ -14,10 +14,14 @@ use std::collections::HashMap; use std::collections::hash_map::Entry; +use std::fmt::Write; use std::future::Future; +use std::ops::Range; use std::sync::Arc; +use std::sync::LazyLock; use databend_common_base::runtime::spawn; +use databend_common_catalog::plan::PartInfoPtr; use databend_common_catalog::plan::Projection; use databend_common_catalog::plan::block_id_in_segment; use databend_common_catalog::plan::block_idx_in_segment; @@ -36,6 +40,7 @@ use databend_storages_common_cache::CacheValue; use databend_storages_common_cache::InMemoryLruCache; use databend_storages_common_cache::LoadParams; use databend_storages_common_io::ReadSettings; +use databend_storages_common_pruner::BlockMetaIndex; use databend_storages_common_table_meta::meta::BlockMeta; use databend_storages_common_table_meta::meta::ColumnMeta; use databend_storages_common_table_meta::meta::Compression; @@ -48,7 +53,9 @@ use tokio::task::JoinHandle; use super::fuse_rows_fetcher::RowsFetchMetadata; use super::fuse_rows_fetcher::RowsFetcher; +use super::read_block_context::read_parquet_page_range_data; use crate::BlockReadResult; +use crate::FUSE_OPT_KEY_DATA_PAGE_ROWS; use crate::FuseBlockPartInfo; use crate::FuseTable; use crate::io::BlockReader; @@ -95,10 +102,20 @@ pub struct RowsFetchMetadataImpl { pub location: String, pub nums_rows: usize, + pub file_size: u64, pub compression: Compression, pub columns_meta: HashMap, + pub block_meta_index: Option, } +static ROW_FETCH_METADATA_CACHE: LazyLock> = + LazyLock::new(|| { + InMemoryLruCache::with_items_capacity( + String::from("RowFetchProjectedBlockMetaCache"), + 16_384, + ) + }); + impl RowsFetchMetadata for RowsFetchMetadataImpl { fn row_bytes(&self) -> usize { self.row_bytes @@ -127,6 +144,8 @@ pub(super) struct ParquetRowsFetcher { segment_reader: CompactSegmentInfoReader, block_meta_lru_cache: InMemoryLruCache, + source_block_meta_indexes: HashMap, + metadata_cache_prefix: String, } #[async_trait::async_trait] @@ -135,16 +154,22 @@ impl RowsFetcher for ParquetRowsFetcher { #[async_backtrace::framed] async fn initialize(&mut self) -> Result<()> { - self.snapshot = self.table.read_table_snapshot().await?; Ok(()) } async fn fetch_metadata(&mut self, block_id: u64) -> Result { + let global_key = self.metadata_cache_key(block_id); + if let Some(metadata) = ROW_FETCH_METADATA_CACHE.get(&global_key) { + return Ok(metadata); + } if let Some(v) = self.block_meta_lru_cache.get(block_id.to_string()) { return Ok(v.clone()); } // load metadata let (segment, block) = split_prefix(block_id); + if self.snapshot.is_none() { + self.snapshot = self.table.read_table_snapshot().await?; + } let snapshot = self.snapshot.as_ref().unwrap(); let (location, ver) = snapshot.segments[segment as usize].clone(); @@ -166,18 +191,22 @@ impl RowsFetcher for ParquetRowsFetcher { cache_end = std::cmp::min(cache_end, blocks.len()); let metadata = self.build_metadata(&blocks[cache_start..cache_end])?; - for (block_index, metadata) in (cache_start..cache_end).zip(metadata.into_iter()) { + for (block_index, mut metadata) in (cache_start..cache_end).zip(metadata.into_iter()) { let block_id = block_id_in_segment(blocks.len(), block_index); let block_id = compute_row_id_prefix(segment, block_id as u64); - self.block_meta_lru_cache - .insert(block_id.to_string(), metadata); + metadata.block_meta_index = self.source_block_meta_indexes.get(&block_id).cloned(); + if metadata.block_meta_index.is_some() { + ROW_FETCH_METADATA_CACHE.insert(self.metadata_cache_key(block_id), metadata); + } else { + self.block_meta_lru_cache + .insert(block_id.to_string(), metadata); + } } - Ok(self - .block_meta_lru_cache - .get(block_id.to_string()) - .clone() - .unwrap()) + ROW_FETCH_METADATA_CACHE + .get(global_key) + .or_else(|| self.block_meta_lru_cache.get(block_id.to_string())) + .ok_or_else(|| ErrorCode::Internal("Row fetch block metadata is missing")) } #[async_backtrace::framed] @@ -288,10 +317,25 @@ impl ParquetRowsFetcher { settings: ReadSettings, max_threads: usize, io_semaphore: Arc, + source_parts: &[PartInfoPtr], ) -> Self { let schema = table.schema(); let operator = table.operator.clone(); let segment_reader = MetaReaders::segment_info_reader(operator, schema.clone()); + let metadata_cache_prefix = projection_cache_prefix(&table, &projection); + let data_page_rows = table.get_option(FUSE_OPT_KEY_DATA_PAGE_ROWS, 0usize); + let source_block_meta_indexes = source_parts + .iter() + .filter_map(|part| FuseBlockPartInfo::from_part(part).ok()) + .filter_map(|part| part.block_meta_index.clone()) + .map(|mut index| { + if data_page_rows > 0 { + index.page_size = data_page_rows; + } + let prefix = compute_row_id_prefix(index.segment_idx as u64, index.block_id as u64); + (prefix, index) + }) + .collect(); ParquetRowsFetcher { table, snapshot: None, @@ -306,9 +350,15 @@ impl ParquetRowsFetcher { String::from("RowFetchBlockMetaCache"), 128, ), + source_block_meta_indexes, + metadata_cache_prefix, } } + fn metadata_cache_key(&self, block_id: u64) -> String { + format!("{}{}", self.metadata_cache_prefix, block_id) + } + fn fetch_block( &self, metadata: Arc, @@ -327,6 +377,13 @@ impl ParquetRowsFetcher { .acquire_owned() .await .expect("row-fetch io semaphore never closed"); + if let Some(block) = + Self::fetch_page_range(&reader, &settings, &metadata, take_indices.as_slice()) + .await? + { + return Ok((final_index, block)); + } + let chunk = reader .read_columns_data_by_merge_io( &settings, @@ -344,6 +401,47 @@ impl ParquetRowsFetcher { } } + async fn fetch_page_range( + reader: &Arc, + settings: &ReadSettings, + metadata: &RowsFetchMetadataImpl, + take_indices: &[u32], + ) -> Result> { + let Some(mut block_meta_index) = metadata.block_meta_index.clone() else { + return Ok(None); + }; + let Some((range, range_start, relative_indices)) = + page_range_for_rows(take_indices, metadata.nums_rows, block_meta_index.page_size)? + else { + return Ok(None); + }; + + block_meta_index.range = Some(range); + let part = FuseBlockPartInfo { + location: metadata.location.clone(), + bloom_filter_index_location: None, + bloom_filter_index_size: 0, + create_on: None, + nums_rows: metadata.nums_rows, + file_size: metadata.file_size, + columns_meta: metadata.columns_meta.clone(), + columns_stat: None, + compression: metadata.compression, + sort_min_max: None, + block_meta_index: Some(block_meta_index), + }; + + let Some(block) = read_parquet_page_range_data(reader, settings, &part).await? else { + return Ok(None); + }; + debug_assert!( + relative_indices + .iter() + .all(|index| *index as usize + range_start < metadata.nums_rows) + ); + Ok(Some(block.take(relative_indices.as_slice())?)) + } + fn build_metadata(&self, meta: &[Arc]) -> Result> { let arrow_schema = self.schema.as_ref().into(); let column_nodes = ColumnNodes::new_from_schema(&arrow_schema, Some(&self.schema)); @@ -383,9 +481,11 @@ impl ParquetRowsFetcher { row_bytes: average_bytes, block_bytes, nums_rows: fuse_part.nums_rows, + file_size: block_meta.file_size, compression: fuse_part.compression, location: fuse_part.location.clone(), columns_meta: fuse_part.columns_meta.clone(), + block_meta_index: None, }); } @@ -409,6 +509,86 @@ impl ParquetRowsFetcher { } } +fn projection_cache_prefix(table: &FuseTable, projection: &Projection) -> String { + let table_ident = table.get_table_info().ident; + metadata_cache_prefix( + &table.get_operator_ref().info().root(), + table_ident.table_id, + table_ident.seq, + &table.query_result_cache_id(), + table.get_option(FUSE_OPT_KEY_DATA_PAGE_ROWS, 0usize), + projection, + ) +} + +fn metadata_cache_prefix( + operator_root: &str, + table_id: u64, + table_seq: u64, + snapshot_cache_id: &str, + data_page_rows: usize, + projection: &Projection, +) -> String { + let mut key = format!( + "row-fetch-meta-v3|{}:{}|{}|{}|{}:{}|{}|", + operator_root.len(), + operator_root, + table_id, + table_seq, + snapshot_cache_id.len(), + snapshot_cache_id, + data_page_rows, + ); + match projection { + Projection::Columns(indices) => { + key.push('c'); + for index in indices { + write!(key, ":{index}").unwrap(); + } + } + Projection::InnerColumns(paths) => { + key.push('i'); + for (index, path) in paths { + write!(key, ":{index}=").unwrap(); + for child in path { + write!(key, "{child},").unwrap(); + } + } + } + } + key.push(':'); + key +} + +type PageRangeSelection = (Range, usize, Vec); + +fn page_range_for_rows( + take_indices: &[u32], + nums_rows: usize, + page_size: usize, +) -> Result> { + if take_indices.is_empty() || page_size == 0 || page_size >= nums_rows { + return Ok(None); + } + + let min_row = *take_indices.iter().min().unwrap() as usize; + let max_row = *take_indices.iter().max().unwrap() as usize; + if max_row >= nums_rows { + return Err(ErrorCode::Internal(format!( + "RowID index {max_row} is outside a block with {nums_rows} rows" + ))); + } + + let start_page = min_row / page_size; + let end_page = max_row / page_size + 1; + let range_start = start_page * page_size; + let relative_indices = take_indices + .iter() + .map(|index| *index - range_start as u32) + .collect(); + Ok(Some((start_page..end_page, range_start, relative_indices))) +} + impl From for CacheValue { fn from(value: RowsFetchMetadataImpl) -> Self { CacheValue::new(value, 0) @@ -481,4 +661,59 @@ mod tests { drop(permit); assert_eq!(semaphore.available_permits(), 1); } + + #[test] + fn test_page_range_for_rows_preserves_output_order() { + let (range, start, relative) = page_range_for_rows(&[4100, 4097, 8192, 4100], 12_000, 4096) + .unwrap() + .unwrap(); + + assert_eq!(range, 1..3); + assert_eq!(start, 4096); + assert_eq!(relative, vec![4, 1, 4096, 4]); + } + + #[test] + fn test_page_range_for_rows_rejects_invalid_row_id() { + let err = page_range_for_rows(&[10], 10, 4).unwrap_err(); + assert!(err.message().contains("outside a block")); + } + + #[test] + fn test_metadata_cache_prefix_isolates_storage_and_snapshot() { + let projection = Projection::Columns(vec![1, 3]); + let base = metadata_cache_prefix("root-a", 1, 2, "snapshot-a", 100, &projection); + + assert_ne!( + base, + metadata_cache_prefix("root-b", 1, 2, "snapshot-a", 100, &projection) + ); + assert_ne!( + base, + metadata_cache_prefix("root-a", 1, 2, "snapshot-b", 100, &projection) + ); + assert_ne!( + base, + metadata_cache_prefix("root-a", 2, 2, "snapshot-a", 100, &projection) + ); + assert_ne!( + base, + metadata_cache_prefix("root-a", 1, 3, "snapshot-a", 100, &projection) + ); + assert_ne!( + base, + metadata_cache_prefix("root-a", 1, 2, "snapshot-a", 200, &projection) + ); + assert_ne!( + metadata_cache_prefix("ab", 1, 2, "c", 100, &projection), + metadata_cache_prefix("a", 1, 2, "bc", 100, &projection) + ); + } + + #[test] + fn test_page_range_for_rows_skips_empty_or_unknown_pages() { + assert!(page_range_for_rows(&[], 10, 4).unwrap().is_none()); + assert!(page_range_for_rows(&[1], 10, 0).unwrap().is_none()); + assert!(page_range_for_rows(&[1], 10, 10).unwrap().is_none()); + } } diff --git a/src/query/storages/fuse/src/operations/read/read_block_context.rs b/src/query/storages/fuse/src/operations/read/read_block_context.rs index a0121ee638a9e..c805575f30c7c 100644 --- a/src/query/storages/fuse/src/operations/read/read_block_context.rs +++ b/src/query/storages/fuse/src/operations/read/read_block_context.rs @@ -14,11 +14,25 @@ use std::sync::Arc; +use arrow_schema::Schema; use databend_common_catalog::plan::PartInfoPtr; use databend_common_catalog::table_context::TableContext; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_expression::DataBlock; +use databend_common_storages_parquet::InMemoryRowGroup; +use databend_common_storages_parquet::ParquetFileReader; +use databend_common_storages_parquet::ReadSettings as ParquetReadSettings; +use databend_storages_common_cache::CacheAccessor; +use databend_storages_common_cache::CacheManager; use databend_storages_common_io::ReadSettings; use log::debug; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::ParquetRecordBatchReader; +use parquet::arrow::parquet_to_arrow_field_levels; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::metadata::ParquetMetaDataReader; use super::block_format::FuseParquetBlockFormat; use super::parquet_data_source::ParquetDataSource; @@ -26,6 +40,8 @@ use crate::FuseBlockPartInfo; use crate::FuseStorageFormat; use crate::io::AggIndexReader; use crate::io::BlockReadContext; +use crate::io::BlockReader; +use crate::io::RowSelection; use crate::io::TableMetaLocationGenerator; use crate::io::VirtualBlockReadResult; use crate::io::VirtualColumnReader; @@ -168,3 +184,117 @@ impl ReadBlockContext { .await } } + +pub(crate) async fn read_parquet_page_range_data( + block_reader: &Arc, + read_settings: &ReadSettings, + part: &FuseBlockPartInfo, +) -> Result> { + let Some(page_cache_key) = block_reader.page_range_data_cache_key(part) else { + return Ok(None); + }; + if let Some(data_block) = block_reader.cached_page_range_data(&page_cache_key) { + return Ok(Some(data_block)); + } + + let block_read_ctx = block_reader.read_context(); + let metadata = parquet_metadata_with_offset_indexes(&block_read_ctx, part).await?; + if metadata.num_row_groups() != 1 { + return Ok(None); + } + let Some(offset_indexes) = metadata.offset_index().and_then(|v| v.first()) else { + return Ok(None); + }; + + let schema_descr = metadata.file_metadata().schema_descr(); + let projection_indices = block_read_ctx + .project_indices() + .iter() + .filter_map(|(index, (column_id, ..))| { + part.columns_meta.contains_key(column_id).then_some(*index) + }) + .collect::>(); + if projection_indices.is_empty() + || projection_indices + .iter() + .any(|index| *index >= schema_descr.num_columns()) + { + return Ok(None); + } + + let page_bitmap = BlockReader::page_range_bitmap(part) + .ok_or_else(|| ErrorCode::Internal("page range is missing"))?; + let parquet_selection = RowSelection::from(&page_bitmap).selection; + let page_locations = offset_indexes + .iter() + .map(|index| index.page_locations().to_vec()) + .collect::>(); + if page_locations.len() != schema_descr.num_columns() { + return Ok(None); + } + + let projection = ProjectionMask::leaves(schema_descr, projection_indices); + let parquet_read_settings = ParquetReadSettings { + max_gap_size: read_settings.max_gap_size, + max_range_size: read_settings.max_range_size, + parquet_fast_read_bytes: read_settings.parquet_fast_read_bytes, + enable_cache: true, + }; + let mut row_group = InMemoryRowGroup::new( + &part.location, + block_read_ctx.operator().clone(), + metadata.row_group(0), + Some(page_locations), + parquet_read_settings, + ); + row_group + .fetch(&projection, Some(&parquet_selection)) + .await?; + + let arrow_schema = Schema::from(block_reader.original_schema.as_ref()); + let field_levels = + parquet_to_arrow_field_levels(schema_descr, projection, Some(arrow_schema.fields()))?; + let mut reader = ParquetRecordBatchReader::try_new_with_row_groups( + &field_levels, + &row_group, + part.nums_rows, + Some(parquet_selection), + )?; + let record_batch = reader + .next() + .ok_or_else(|| ErrorCode::Internal("selected parquet range returned no rows"))??; + debug_assert!(reader.next().is_none()); + + let data_block = block_reader.deserialize_parquet_record_batch(part, &record_batch)?; + Ok(Some( + block_reader.cache_page_range_data(page_cache_key, data_block), + )) +} + +async fn parquet_metadata_with_offset_indexes( + block_read_ctx: &BlockReadContext, + part: &FuseBlockPartInfo, +) -> Result> { + let cache = CacheManager::instance().get_parquet_meta_data_cache(); + let cache_key = format!( + "{}{}", + block_read_ctx.operator().info().root(), + part.location + ); + if let Some(metadata) = cache.as_ref().and_then(|cache| cache.get(&cache_key)) { + if metadata.offset_index().is_some() { + return Ok(metadata); + } + } + + let op_reader = block_read_ctx.operator().reader(&part.location).await?; + let mut file_reader = ParquetFileReader::new(op_reader, part.file_size); + let metadata = ParquetMetaDataReader::new() + .with_offset_index_policy(PageIndexPolicy::Optional) + .load_and_finish(&mut file_reader, part.file_size) + .await?; + Ok(match cache { + Some(cache) => cache.insert(cache_key, metadata), + None => Arc::new(metadata), + }) +} diff --git a/src/query/storages/system/src/caches_table.rs b/src/query/storages/system/src/caches_table.rs index d4245244ca728..8bdc5ea652978 100644 --- a/src/query/storages/system/src/caches_table.rs +++ b/src/query/storages/system/src/caches_table.rs @@ -93,7 +93,7 @@ impl SyncSystemTable for CachesTable { let prune_partitions_cache = cache_manager.get_prune_partitions_cache(); let parquet_meta_data_cache = cache_manager.get_parquet_meta_data_cache(); let column_data_cache = cache_manager.get_column_data_cache(); - let table_column_array_cache = cache_manager.get_table_data_array_cache(); + let table_data_cache = cache_manager.get_table_data_cache(); let iceberg_table_cache = cache_manager.get_iceberg_table_cache(); let mut columns = CachesTableColumns::default(); @@ -189,8 +189,8 @@ impl SyncSystemTable for CachesTable { Self::append_row(&parquet_meta_data_cache, &local_node, &mut columns); } - if let Some(table_column_array_cache) = table_column_array_cache { - Self::append_row(&table_column_array_cache, &local_node, &mut columns); + if let Some(table_data_cache) = table_data_cache { + Self::append_row(&table_data_cache, &local_node, &mut columns); } if let Some(iceberg_table_cache) = iceberg_table_cache { diff --git a/tests/task/test-private-task.sh b/tests/task/test-private-task.sh index 92ab0a19ad054..74bd62383a48e 100644 --- a/tests/task/test-private-task.sh +++ b/tests/task/test-private-task.sh @@ -706,6 +706,26 @@ else exit 1 fi +actual_child_success_count=0 +for _ in {1..20}; do + response=$(query_sql_with_auth "root:" "SELECT count(*) FROM system_task.task_run WHERE task_name IN ('fanout_child_a', 'fanout_child_b', 'fanout_child_c') AND state = 'SUCCEEDED' AND completed_at IS NOT NULL") + check_response_error "$response" + actual_child_success_count=$(echo "$response" | jq -r '.data[0][0]') + if [ "$actual_child_success_count" = "3" ]; then + break + fi + sleep 1 +done + +if [ "$actual_child_success_count" = "3" ]; then + echo "✅ Private task fan-out child runs are terminal" +else + echo "❌ Expected all private task fan-out child runs to be terminal" + echo "Expected: 3" + echo "Actual : $actual_child_success_count" + exit 1 +fi + response=$(query_sql_with_auth "root:" "UPDATE system_task.task_run SET state = 'SKIPPED', error_code = 0, error_message = 'OVERLAPPING_EXECUTION: test', completed_at = to_timestamp(4102444800) WHERE task_name = 'fanout_root'") check_response_error "$response"