From 09b368c00aff85a468497c213100a27210f265fc Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Wed, 8 Apr 2026 16:34:13 -0700 Subject: [PATCH 001/120] feat: tune constants (#7851) * feat: tune constants Signed-off-by: Ruihang Xia * cap output batch size Signed-off-by: Ruihang Xia * handle empty input Signed-off-by: Ruihang Xia * one more ut for cr Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- .../datasource/src/file_format/parquet.rs | 37 ++- src/mito2/src/cache/file_cache.rs | 43 ++- src/mito2/src/memtable/bulk.rs | 2 +- src/mito2/src/memtable/partition_tree/data.rs | 7 +- src/mito2/src/sst/parquet.rs | 12 +- src/promql/src/extension_plan/absent.rs | 277 ++++++++++++++---- src/query/src/range_select/plan.rs | 195 +++++++++++- 7 files changed, 490 insertions(+), 83 deletions(-) diff --git a/src/common/datasource/src/file_format/parquet.rs b/src/common/datasource/src/file_format/parquet.rs index c2c14b4680..9c8e8d6ce8 100644 --- a/src/common/datasource/src/file_format/parquet.rs +++ b/src/common/datasource/src/file_format/parquet.rs @@ -23,7 +23,9 @@ use datafusion::error::Result as DatafusionResult; use datafusion::parquet::arrow::async_reader::AsyncFileReader; use datafusion::parquet::arrow::{ArrowWriter, parquet_to_arrow_schema}; use datafusion::parquet::errors::{ParquetError, Result as ParquetResult}; -use datafusion::parquet::file::metadata::ParquetMetaData; +use datafusion::parquet::file::metadata::{ + PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, +}; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_datasource::PartitionedFile; @@ -94,35 +96,40 @@ impl DefaultParquetFileReaderFactory { } impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { - // TODO(weny): Supports [`metadata_size_hint`]. - // The upstream has a implementation supports [`metadata_size_hint`], - // however it coupled with Box. fn create_reader( &self, _partition_index: usize, partitioned_file: PartitionedFile, - _metadata_size_hint: Option, + metadata_size_hint: Option, _metrics: &ExecutionPlanMetricsSet, ) -> DatafusionResult> { let path = partitioned_file.path().to_string(); let object_store = self.object_store.clone(); - Ok(Box::new(LazyParquetFileReader::new(object_store, path))) + Ok(Box::new(LazyParquetFileReader::new( + object_store, + path, + metadata_size_hint, + ))) } } pub struct LazyParquetFileReader { object_store: ObjectStore, reader: Option>, + file_size: Option, + metadata_size_hint: Option, path: String, } impl LazyParquetFileReader { - pub fn new(object_store: ObjectStore, path: String) -> Self { + pub fn new(object_store: ObjectStore, path: String, metadata_size_hint: Option) -> Self { LazyParquetFileReader { object_store, path, reader: None, + file_size: None, + metadata_size_hint, } } @@ -130,6 +137,7 @@ impl LazyParquetFileReader { async fn maybe_initialize(&mut self) -> result::Result<(), object_store::Error> { if self.reader.is_none() { let meta = self.object_store.stat(&self.path).await?; + self.file_size = Some(meta.content_length()); let reader = self .object_store .reader(&self.path) @@ -166,8 +174,19 @@ impl AsyncFileReader for LazyParquetFileReader { self.maybe_initialize() .await .map_err(|e| ParquetError::External(Box::new(e)))?; - // Safety: Must initialized - self.reader.as_mut().unwrap().get_metadata(options).await + + let metadata_opts = options.map(|o| o.metadata_options().clone()); + let metadata_reader = ParquetMetaDataReader::new() + .with_metadata_options(metadata_opts) + .with_page_index_policy(PageIndexPolicy::from( + options.is_some_and(|o| o.page_index()), + )) + .with_prefetch_hint(self.metadata_size_hint); + + let metadata = metadata_reader + .load_and_finish(self.reader.as_mut().unwrap(), self.file_size.unwrap()) + .await?; + Ok(Arc::new(metadata)) }) } } diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 278838b369..9b987c810b 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -288,6 +288,17 @@ pub(crate) struct FileCache { pub(crate) type FileCacheRef = Arc; impl FileCache { + /// Splits the configured total capacity between parquet and puffin caches + /// without exceeding the requested overall budget. + fn split_cache_capacities(total_capacity: u64, index_percent: u8) -> (u64, u64) { + let desired_puffin_capacity = total_capacity * u64::from(index_percent) / 100; + let min_cache_capacity = MIN_CACHE_CAPACITY.min(total_capacity / 2); + let puffin_capacity = + desired_puffin_capacity.clamp(min_cache_capacity, total_capacity - min_cache_capacity); + let parquet_capacity = total_capacity - puffin_capacity; + (parquet_capacity, puffin_capacity) + } + /// Creates a new file cache. pub(crate) fn new( local_store: ObjectStore, @@ -302,14 +313,8 @@ impl FileCache { .unwrap_or(DEFAULT_INDEX_CACHE_PERCENT); let total_capacity = capacity.as_bytes(); - // Convert percent to ratio and calculate capacity for each cache - let index_ratio = index_percent as f64 / 100.0; - let puffin_capacity = (total_capacity as f64 * index_ratio) as u64; - let parquet_capacity = total_capacity - puffin_capacity; - - // Ensure both capacities are at least 512MB - let puffin_capacity = puffin_capacity.max(MIN_CACHE_CAPACITY); - let parquet_capacity = parquet_capacity.max(MIN_CACHE_CAPACITY); + let (parquet_capacity, puffin_capacity) = + Self::split_cache_capacities(total_capacity, index_percent); info!( "Initializing file cache with index_percent: {}%, total_capacity: {}, parquet_capacity: {}, puffin_capacity: {}", @@ -1064,6 +1069,28 @@ mod tests { assert_eq!(data, bytes[3].as_ref()); } + #[test] + fn test_file_cache_capacity_respects_total_budget() { + let total_capacity = ReadableSize::mb(256).as_bytes(); + let (parquet_capacity, puffin_capacity) = + FileCache::split_cache_capacities(total_capacity, 20); + + assert_eq!(total_capacity, parquet_capacity + puffin_capacity); + assert_eq!(ReadableSize::mb(128).as_bytes(), parquet_capacity); + assert_eq!(ReadableSize::mb(128).as_bytes(), puffin_capacity); + } + + #[test] + fn test_file_cache_capacity_keeps_split_when_total_allows_it() { + let total_capacity = ReadableSize::gb(5).as_bytes(); + let (parquet_capacity, puffin_capacity) = + FileCache::split_cache_capacities(total_capacity, 20); + + assert_eq!(total_capacity, parquet_capacity + puffin_capacity); + assert_eq!(ReadableSize::gb(4).as_bytes(), parquet_capacity); + assert_eq!(ReadableSize::gb(1).as_bytes(), puffin_capacity); + } + #[test] fn test_cache_file_path() { let file_id = FileId::parse_str("3368731b-a556-42b8-a5df-9c31ce155095").unwrap(); diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 9d25d0c39f..24b2bebaa9 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -137,7 +137,7 @@ struct CollectedParts { /// All parts in a bulk memtable. #[derive(Default)] struct BulkParts { - /// Unordered small parts (< 1024 rows). + /// Unordered small parts. unordered_part: UnorderedPart, /// All parts (raw and encoded). parts: Vec, diff --git a/src/mito2/src/memtable/partition_tree/data.rs b/src/mito2/src/memtable/partition_tree/data.rs index a6d40bdcbf..f6e2a59bec 100644 --- a/src/mito2/src/memtable/partition_tree/data.rs +++ b/src/mito2/src/memtable/partition_tree/data.rs @@ -50,6 +50,7 @@ use crate::memtable::partition_tree::merger::{DataBatchKey, DataNode, DataSource use crate::metrics::{ PARTITION_TREE_DATA_BUFFER_FREEZE_STAGE_ELAPSED, PARTITION_TREE_READ_STAGE_ELAPSED, }; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; const PK_INDEX_COLUMN_NAME: &str = "__pk_index"; @@ -821,7 +822,11 @@ impl DataPart { /// Reads frozen data part and yields [DataBatch]es. pub fn read(&self) -> Result { match self { - DataPart::Parquet(data_bytes) => DataPartReader::new(data_bytes.data.clone(), None), + // Keep encoded memtable scans aligned with mito/DataFusion batch sizing instead of + // parquet-rs's implicit 1024-row default. + DataPart::Parquet(data_bytes) => { + DataPartReader::new(data_bytes.data.clone(), Some(DEFAULT_READ_BATCH_SIZE)) + } } } diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 2ca83ca8cf..2447824ad9 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -41,9 +41,17 @@ pub mod writer; pub const PARQUET_METADATA_KEY: &str = "greptime:metadata"; /// Default batch size to read parquet files. -pub(crate) const DEFAULT_READ_BATCH_SIZE: usize = 1024; +/// +/// This is a runtime-only scan granularity, so we align it with DataFusion's +/// default execution batch size to reduce rebatching and concatenation in the +/// query pipeline. +pub(crate) const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; /// Default row group size for parquet files. -pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * DEFAULT_READ_BATCH_SIZE; +/// +/// Keep the existing persisted/on-disk default stable. It intentionally stays +/// decoupled from [`DEFAULT_READ_BATCH_SIZE`] so we can tune runtime scan +/// batching without changing the row group layout of newly written SSTs. +pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; /// Parquet write options. #[derive(Debug, Clone)] diff --git a/src/promql/src/extension_plan/absent.rs b/src/promql/src/extension_plan/absent.rs index db31a3d901..71af413029 100644 --- a/src/promql/src/extension_plan/absent.rs +++ b/src/promql/src/extension_plan/absent.rs @@ -49,9 +49,6 @@ use snafu::ResultExt; use crate::error::DeserializeSnafu; use crate::extension_plan::{Millisecond, resolve_column_name, serialize_column_index}; -/// Maximum number of rows per output batch -const ABSENT_BATCH_SIZE: usize = 8192; - #[derive(Debug, PartialEq, Eq, Hash)] pub struct Absent { start: Millisecond, @@ -390,11 +387,13 @@ impl ExecutionPlan for AbsentExec { context: Arc, ) -> DataFusionResult { let baseline_metric = BaselineMetrics::new(&self.metric, partition); + let batch_size = context.session_config().batch_size(); let input = self.input.execute(partition, context)?; Ok(Box::pin(AbsentStream { end: self.end, step: self.step, + batch_size, time_index_column_index: self .input .schema() @@ -407,6 +406,8 @@ impl ExecutionPlan for AbsentExec { metric: baseline_metric, // Buffer for streaming output timestamps output_timestamps: Vec::new(), + input_timestamps: Vec::new(), + input_timestamp_offset: 0, // Current timestamp in the output range output_ts_cursor: self.start, input_finished: false, @@ -441,6 +442,7 @@ impl DisplayAs for AbsentExec { pub struct AbsentStream { end: Millisecond, step: Millisecond, + batch_size: usize, time_index_column_index: usize, output_schema: SchemaRef, fake_labels: Vec<(String, String)>, @@ -448,6 +450,9 @@ pub struct AbsentStream { metric: BaselineMetrics, // Buffer for streaming output timestamps output_timestamps: Vec, + // Current input timestamps being processed incrementally. + input_timestamps: Vec, + input_timestamp_offset: usize, // Current timestamp in the output range output_ts_cursor: Millisecond, input_finished: bool, @@ -464,52 +469,53 @@ impl Stream for AbsentStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - if !self.input_finished { - match ready!(self.input.poll_next_unpin(cx)) { - Some(Ok(batch)) => { - let timer = std::time::Instant::now(); - if let Err(e) = self.process_input_batch(&batch) { - return Poll::Ready(Some(Err(e))); - } - self.metric.elapsed_compute().add_elapsed(timer); - - // If we have enough data for a batch, output it - if self.output_timestamps.len() >= ABSENT_BATCH_SIZE { - let timer = std::time::Instant::now(); - let result = self.flush_output_batch(); - self.metric.elapsed_compute().add_elapsed(timer); - - match result { - Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), - Ok(None) => continue, - Err(e) => return Poll::Ready(Some(Err(e))), - } - } - } - Some(Err(e)) => return Poll::Ready(Some(Err(e))), - None => { - self.input_finished = true; - - let timer = std::time::Instant::now(); - // Process any remaining absent timestamps - if let Err(e) = self.process_remaining_absent_timestamps() { - return Poll::Ready(Some(Err(e))); - } - let result = self.flush_output_batch(); - self.metric.elapsed_compute().add_elapsed(timer); - return Poll::Ready(result.transpose()); - } + if self.has_pending_input_timestamps() { + let timer = std::time::Instant::now(); + if let Err(e) = self.process_input_batch() { + return Poll::Ready(Some(Err(e))); + } + self.metric.elapsed_compute().add_elapsed(timer); + + match self.flush_output_batch() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => continue, + Err(e) => return Poll::Ready(Some(Err(e))), + } + } + + if self.input_finished { + let timer = std::time::Instant::now(); + if let Err(e) = self.process_remaining_absent_timestamps() { + return Poll::Ready(Some(Err(e))); + } + self.metric.elapsed_compute().add_elapsed(timer); + + match self.flush_output_batch() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => return Poll::Ready(None), + Err(e) => return Poll::Ready(Some(Err(e))), + } + } + + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + let timer = std::time::Instant::now(); + if let Err(e) = self.buffer_input_timestamps(&batch) { + return Poll::Ready(Some(Err(e))); + } + self.metric.elapsed_compute().add_elapsed(timer); + } + Some(Err(e)) => return Poll::Ready(Some(Err(e))), + None => { + self.input_finished = true; } - } else { - return Poll::Ready(None); } } } } impl AbsentStream { - fn process_input_batch(&mut self, batch: &RecordBatch) -> DataFusionResult<()> { - // Extract timestamps from this batch + fn buffer_input_timestamps(&mut self, batch: &RecordBatch) -> DataFusionResult<()> { let timestamp_array = batch.column(self.time_index_column_index); let milli_ts_array = arrow::compute::cast( timestamp_array, @@ -519,29 +525,52 @@ impl AbsentStream { .as_any() .downcast_ref::() .unwrap(); + self.input_timestamps.clear(); + self.input_timestamps + .extend_from_slice(timestamp_array.values()); + self.input_timestamp_offset = 0; + Ok(()) + } + + fn has_pending_input_timestamps(&self) -> bool { + self.input_timestamp_offset < self.input_timestamps.len() + } + + fn process_input_batch(&mut self) -> DataFusionResult<()> { + while self.input_timestamp_offset < self.input_timestamps.len() { + let input_ts = self.input_timestamps[self.input_timestamp_offset]; - // Process against current output cursor position - for &input_ts in timestamp_array.values() { // Generate absent timestamps up to this input timestamp while self.output_ts_cursor < input_ts && self.output_ts_cursor <= self.end { self.output_timestamps.push(self.output_ts_cursor); self.output_ts_cursor += self.step; + + if self.output_timestamps.len() >= self.batch_size { + return Ok(()); + } } // Skip the input timestamp if it matches our cursor if self.output_ts_cursor == input_ts { self.output_ts_cursor += self.step; } + + self.input_timestamp_offset += 1; } + self.input_timestamps.clear(); + self.input_timestamp_offset = 0; Ok(()) } fn process_remaining_absent_timestamps(&mut self) -> DataFusionResult<()> { - // Generate all remaining absent timestamps (input is finished) while self.output_ts_cursor <= self.end { self.output_timestamps.push(self.output_ts_cursor); self.output_ts_cursor += self.step; + + if self.output_timestamps.len() >= self.batch_size { + return Ok(()); + } } Ok(()) } @@ -551,11 +580,16 @@ impl AbsentStream { return Ok(None); } + let timestamps = if self.output_timestamps.len() <= self.batch_size { + std::mem::take(&mut self.output_timestamps) + } else { + let remaining = self.output_timestamps.split_off(self.batch_size); + std::mem::replace(&mut self.output_timestamps, remaining) + }; + let mut columns: Vec = Vec::with_capacity(self.output_schema.fields().len()); - let num_rows = self.output_timestamps.len(); - columns.push(Arc::new(TimestampMillisecondArray::from( - self.output_timestamps.clone(), - )) as _); + let num_rows = timestamps.len(); + columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)) as _); columns.push(Arc::new(Float64Array::from(vec![1.0; num_rows])) as _); for (_, value) in self.fake_labels.iter() { @@ -567,7 +601,6 @@ impl AbsentStream { let batch = RecordBatch::try_new(self.output_schema.clone(), columns)?; - self.output_timestamps.clear(); Ok(Some(batch)) } } @@ -580,7 +613,7 @@ mod tests { use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::memory::DataSourceExec; use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::prelude::SessionContext; + use datafusion::prelude::{SessionConfig, SessionContext}; use datatypes::arrow::array::{Float64Array, TimestampMillisecondArray}; use super::*; @@ -725,4 +758,146 @@ mod tests { // Should output all timestamps in range: 0, 1000, 2000 assert_eq!(output_timestamps, vec![0, 1000, 2000]); } + + #[tokio::test] + async fn test_absent_respects_session_batch_size_for_large_gap() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new("value", DataType::Float64, true), + ])); + + let timestamp_array = Arc::new(TimestampMillisecondArray::from(vec![9])); + let value_array = Arc::new(Float64Array::from(vec![1.0])); + let batch = + RecordBatch::try_new(schema.clone(), vec![timestamp_array, value_array]).unwrap(); + + let memory_exec = DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + )); + + let output_schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new("value", DataType::Float64, true), + ])); + + let absent_exec = AbsentExec { + start: 0, + end: 10, + step: 1, + time_index_column: "timestamp".to_string(), + value_column: "value".to_string(), + fake_labels: vec![], + output_schema: output_schema.clone(), + input: Arc::new(memory_exec), + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + metric: ExecutionPlanMetricsSet::new(), + }; + + let session_ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(3)); + let task_ctx = session_ctx.task_ctx(); + let mut stream = absent_exec.execute(0, task_ctx).unwrap(); + + let mut batch_sizes = Vec::new(); + let mut output_timestamps = Vec::new(); + while let Some(batch_result) = stream.next().await { + let batch = batch_result.unwrap(); + batch_sizes.push(batch.num_rows()); + + let ts_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..ts_array.len() { + if !ts_array.is_null(i) { + output_timestamps.push(ts_array.value(i)); + } + } + } + + assert_eq!(batch_sizes, vec![3, 3, 3, 1]); + assert_eq!(output_timestamps, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 10]); + } + + #[tokio::test] + async fn test_absent_resumes_same_input_timestamp_after_batch_flush() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new("value", DataType::Float64, true), + ])); + + let timestamp_array = Arc::new(TimestampMillisecondArray::from(vec![9])); + let value_array = Arc::new(Float64Array::from(vec![1.0])); + let batch = + RecordBatch::try_new(schema.clone(), vec![timestamp_array, value_array]).unwrap(); + + let memory_exec = DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + )); + + let output_schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new("value", DataType::Float64, true), + ])); + + let absent_exec = AbsentExec { + start: 0, + end: 9, + step: 1, + time_index_column: "timestamp".to_string(), + value_column: "value".to_string(), + fake_labels: vec![], + output_schema: output_schema.clone(), + input: Arc::new(memory_exec), + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + metric: ExecutionPlanMetricsSet::new(), + }; + + let session_ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(3)); + let task_ctx = session_ctx.task_ctx(); + let mut stream = absent_exec.execute(0, task_ctx).unwrap(); + + let mut output_timestamps = Vec::new(); + while let Some(batch_result) = stream.next().await { + let batch = batch_result.unwrap(); + let ts_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..ts_array.len() { + if !ts_array.is_null(i) { + output_timestamps.push(ts_array.value(i)); + } + } + } + + assert_eq!(output_timestamps, vec![0, 1, 2, 3, 4, 5, 6, 7, 8]); + } } diff --git a/src/query/src/range_select/plan.rs b/src/query/src/range_select/plan.rs index e863aaced0..d83514f435 100644 --- a/src/query/src/range_select/plan.rs +++ b/src/query/src/range_select/plan.rs @@ -836,6 +836,7 @@ impl ExecutionPlan for RangeSelectExec { context: Arc, ) -> DfResult { let baseline_metric = BaselineMetrics::new(&self.metric, partition); + let batch_size = context.session_config().batch_size(); let input = self.input.execute(partition, context)?; let schema = input.schema(); let time_index = schema @@ -852,6 +853,7 @@ impl ExecutionPlan for RangeSelectExec { .collect(), )?; Ok(Box::pin(RangeSelectStream { + batch_size, schema: self.schema.clone(), range_exec: self.range_exec.clone(), input, @@ -868,6 +870,8 @@ impl ExecutionPlan for RangeSelectExec { metric: baseline_metric, schema_project: self.schema_project.clone(), schema_before_project: self.schema_before_project.clone(), + output_batch: None, + output_batch_offset: 0, })) } @@ -881,6 +885,7 @@ impl ExecutionPlan for RangeSelectExec { } struct RangeSelectStream { + batch_size: usize, /// the schema of output column schema: SchemaRef, range_exec: Vec, @@ -907,6 +912,8 @@ struct RangeSelectStream { metric: BaselineMetrics, schema_project: Option>, schema_before_project: SchemaRef, + output_batch: Option, + output_batch_offset: usize, } #[derive(Debug)] @@ -1149,6 +1156,36 @@ impl RangeSelectStream { }; Ok(project_output) } + + fn next_output_batch(&mut self) -> DfResult> { + if self.output_batch.is_none() { + self.output_batch = Some(self.generate_output()?); + self.output_batch_offset = 0; + } + + let num_rows = self.output_batch.as_ref().unwrap().num_rows(); + if num_rows == 0 { + self.output_batch = None; + self.output_batch_offset = 0; + return Ok(None); + } + + if self.output_batch_offset == 0 && num_rows <= self.batch_size { + return Ok(self.output_batch.take()); + } + + let offset = self.output_batch_offset; + let len = (num_rows - offset).min(self.batch_size); + let batch = self.output_batch.as_ref().unwrap().slice(offset, len); + self.output_batch_offset += len; + + if self.output_batch_offset >= num_rows { + self.output_batch = None; + self.output_batch_offset = 0; + } + + Ok(Some(batch)) + } } enum ExecutionState { @@ -1191,13 +1228,19 @@ impl Stream for RangeSelectStream { } } ExecutionState::ProducingOutput => { - let result = self.generate_output(); + let result = self.next_output_batch(); return match result { // made output - Ok(batch) => { - self.exec_state = ExecutionState::Done; + Ok(Some(batch)) => { + if self.output_batch.is_none() { + self.exec_state = ExecutionState::Done; + } Poll::Ready(Some(Ok(batch))) } + Ok(None) => { + self.exec_state = ExecutionState::Done; + Poll::Ready(None) + } // error making output Err(error) => Poll::Ready(Some(Err(error))), }; @@ -1251,7 +1294,7 @@ mod test { use datafusion::prelude::SessionContext; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; - use datatypes::arrow::array::TimestampMillisecondArray; + use datatypes::arrow::array::{Float64Array, Int64Array, TimestampMillisecondArray}; use datatypes::arrow_array::StringArray; use super::*; @@ -1313,15 +1356,49 @@ mod test { )) } - async fn do_range_select_test( + fn prepare_empty_test_data(is_float: bool) -> DataSourceExec { + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), + Field::new( + "value", + if is_float { + DataType::Float64 + } else { + DataType::Int64 + }, + true, + ), + Field::new("host", DataType::Utf8, true), + ])); + let timestamp_column: Arc = + Arc::new(TimestampMillisecondArray::from(Vec::::new())) as _; + let value_column: Arc = if is_float { + Arc::new(Float64Array::from(Vec::>::new())) as _ + } else { + Arc::new(Int64Array::from(Vec::>::new())) as _ + }; + let host_column: Arc = + Arc::new(StringArray::from(Vec::>::new())) as _; + let data = RecordBatch::try_new( + schema.clone(), + vec![timestamp_column, value_column, host_column], + ) + .unwrap(); + + DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![data]], schema, None).unwrap(), + )) + } + + async fn collect_range_select_test( range1: Millisecond, range2: Millisecond, align: Millisecond, fill: Option, is_float: bool, is_gap: bool, - expected: String, - ) { + batch_size: usize, + ) -> Vec { let data_type = if is_float { DataType::Float64 } else { @@ -1412,11 +1489,25 @@ mod test { .into(), range_select_exec, ); - let session_context = SessionContext::default(); + let session_context = SessionContext::new_with_config( + datafusion::execution::config::SessionConfig::new().with_batch_size(batch_size), + ); + datafusion::physical_plan::collect(Arc::new(sort_exec), session_context.task_ctx()) + .await + .unwrap() + } + + async fn do_range_select_test( + range1: Millisecond, + range2: Millisecond, + align: Millisecond, + fill: Option, + is_float: bool, + is_gap: bool, + expected: String, + ) { let result = - datafusion::physical_plan::collect(Arc::new(sort_exec), session_context.task_ctx()) - .await - .unwrap(); + collect_range_select_test(range1, range2, align, fill, is_float, is_gap, 8192).await; let result_literal = arrow::util::pretty::pretty_format_batches(&result) .unwrap() @@ -1700,6 +1791,88 @@ mod test { .await; } + #[tokio::test] + async fn range_select_respects_session_batch_size() { + let result = + collect_range_select_test(10_000, 5_000, 5_000, Some(Fill::Null), true, false, 3).await; + + let row_counts = result + .iter() + .map(|batch| batch.num_rows()) + .collect::>(); + assert_eq!(vec![3, 3, 3, 3], row_counts); + } + + #[tokio::test] + async fn range_select_skips_empty_output_batch() { + let memory_exec = Arc::new(prepare_empty_test_data(true)); + let schema = Arc::new(Schema::new(vec![ + Field::new("MIN(value)", DataType::Float64, true), + Field::new("MAX(value)", DataType::Float64, true), + Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), + Field::new("host", DataType::Utf8, true), + ])); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + let input_schema = memory_exec.schema().clone(); + let range_select_exec = Arc::new(RangeSelectExec { + input: memory_exec, + range_exec: vec![ + RangeFnExec { + expr: Arc::new( + AggregateExprBuilder::new( + min_max::min_udaf(), + vec![Arc::new(Column::new("value", 1))], + ) + .schema(input_schema.clone()) + .alias("MIN(value)") + .build() + .unwrap(), + ), + range: 10_000, + fill: Some(Fill::Null), + need_cast: None, + }, + RangeFnExec { + expr: Arc::new( + AggregateExprBuilder::new( + min_max::max_udaf(), + vec![Arc::new(Column::new("value", 1))], + ) + .schema(input_schema) + .alias("MAX(value)") + .build() + .unwrap(), + ), + range: 5_000, + fill: Some(Fill::Null), + need_cast: None, + }, + ], + align: 5_000, + align_to: 0, + by: vec![Arc::new(Column::new("host", 2))], + time_index: TIME_INDEX_COLUMN.to_string(), + schema: schema.clone(), + schema_before_project: schema.clone(), + schema_project: None, + by_schema: Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])), + metric: ExecutionPlanMetricsSet::new(), + cache, + }); + let session_context = SessionContext::new(); + let result = + datafusion::physical_plan::collect(range_select_exec, session_context.task_ctx()) + .await + .unwrap(); + + assert!(result.is_empty()); + } + #[test] fn fill_test() { assert!(Fill::try_from_str("", &DataType::UInt8).unwrap().is_none()); From dca451c4851fb00723122013ab95f5ba8c8301a3 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 9 Apr 2026 11:40:14 +0800 Subject: [PATCH 002/120] fix: remap peer addresses during retries (#7933) * fix: remap peer addresses during retries Signed-off-by: WenyXu * chore: styling Signed-off-by: WenyXu * test: add tests Signed-off-by: WenyXu * chore: apply suggestions from CR Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/ddl/create_table.rs | 27 ++++++-- src/common/meta/src/ddl/drop_database.rs | 6 +- .../meta/src/ddl/drop_database/cursor.rs | 3 + .../meta/src/ddl/drop_database/executor.rs | 69 ++++++++++++++++++- .../meta/src/ddl/drop_database/metadata.rs | 2 + .../meta/src/ddl/drop_database/start.rs | 3 + src/common/meta/src/ddl/drop_table.rs | 26 ++++++- src/common/meta/src/ddl/test_util.rs | 19 +++++ src/common/meta/src/ddl/tests/create_table.rs | 23 ++++++- src/common/meta/src/ddl/tests/drop_table.rs | 48 ++++++++++++- src/common/meta/src/ddl/truncate_table.rs | 19 ++--- src/common/meta/src/ddl_manager.rs | 40 +++++------ src/common/meta/src/key.rs | 2 +- src/common/meta/src/key/table_route.rs | 53 +++++++++----- src/common/procedure/src/local/runner.rs | 54 +++++++++++++++ src/common/procedure/src/procedure.rs | 12 ++++ src/meta-srv/src/procedure/tests.rs | 2 +- 17 files changed, 346 insertions(+), 62 deletions(-) diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index a5b642f1a2..b377a60406 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -172,8 +172,24 @@ impl CreateTableProcedure { /// - [Code::Cancelled](tonic::status::Code::Cancelled) /// - [Code::DeadlineExceeded](tonic::status::Code::DeadlineExceeded) /// - [Code::Unavailable](tonic::status::Code::Unavailable) - pub async fn on_datanode_create_regions(&mut self) -> Result { - let table_route = self.table_route()?.clone(); + pub async fn on_datanode_create_regions(&mut self, retrying: bool) -> Result { + let mut table_route = self.table_route()?.clone(); + if retrying { + info!( + "Remapping region routes addresses for retrying create regions for table: {}", + self.data.table_ref() + ); + let storage = self + .context + .table_metadata_manager + .table_route_manager() + .table_route_storage(); + // The peer addresses may change during retries, + // so we always remap the region routes. + storage + .remap_region_routes(&mut table_route.region_routes) + .await?; + } // Registers opening regions let guards = self.register_opening_regions(&self.context, &table_route.region_routes)?; if !guards.is_empty() { @@ -301,7 +317,10 @@ impl Procedure for CreateTableProcedure { match state { CreateTableState::Prepare => self.on_prepare().await, - CreateTableState::DatanodeCreateRegions => self.on_datanode_create_regions().await, + CreateTableState::DatanodeCreateRegions => { + let retrying = ctx.is_retrying().await.unwrap_or(false); + self.on_datanode_create_regions(retrying).await + } CreateTableState::CreateMetadata => self.on_create_metadata(ctx.procedure_id).await, } .map_err(map_to_procedure_error) @@ -339,7 +358,7 @@ pub struct CreateTableData { #[serde(default)] pub column_metadatas: Vec, /// None stands for not allocated yet. - table_route: Option, + pub(crate) table_route: Option, /// None stands for not allocated yet. pub region_wal_options: Option>, } diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index d3c5c5831e..53b6190ff5 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -58,6 +58,7 @@ pub(crate) struct DropDatabaseContext { schema: String, drop_if_exists: bool, tables: Option>>, + retrying: bool, } #[async_trait::async_trait] @@ -90,6 +91,7 @@ impl DropDatabaseProcedure { schema, drop_if_exists, tables: None, + retrying: false, }, state: Box::new(DropDatabaseStart), } @@ -110,6 +112,7 @@ impl DropDatabaseProcedure { schema, drop_if_exists, tables: None, + retrying: false, }, state, }) @@ -136,9 +139,10 @@ impl Procedure for DropDatabaseProcedure { }) } - async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult { + async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult { let state = &mut self.state; + self.context.retrying = ctx.is_retrying().await.unwrap_or(false); let (next, status) = state .next(&self.runtime_context, &mut self.context) .await diff --git a/src/common/meta/src/ddl/drop_database/cursor.rs b/src/common/meta/src/ddl/drop_database/cursor.rs index 0a1180a52e..e060567be5 100644 --- a/src/common/meta/src/ddl/drop_database/cursor.rs +++ b/src/common/meta/src/ddl/drop_database/cursor.rs @@ -224,6 +224,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; // Ticks let (mut state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); @@ -259,6 +260,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; // Ticks let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); @@ -287,6 +289,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; // Ticks let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); diff --git a/src/common/meta/src/ddl/drop_database/executor.rs b/src/common/meta/src/ddl/drop_database/executor.rs index c478cdb746..80c730f30a 100644 --- a/src/common/meta/src/ddl/drop_database/executor.rs +++ b/src/common/meta/src/ddl/drop_database/executor.rs @@ -96,10 +96,25 @@ impl State for DropDatabaseExecutor { async fn next( &mut self, ddl_ctx: &DdlContext, - _ctx: &mut DropDatabaseContext, + ctx: &mut DropDatabaseContext, ) -> Result<(Box, Status)> { self.register_dropping_regions(ddl_ctx)?; let executor = DropTableExecutor::new(self.table_name.clone(), self.table_id, true); + if ctx.retrying { + info!( + "Remapping region routes addresses for retrying drop regions for table_id: {}", + self.table_id + ); + let storage = ddl_ctx + .table_metadata_manager + .table_route_manager() + .table_route_storage(); + // The peer addresses may change during retries, + // so we always remap the region routes. + storage + .remap_region_routes(&mut self.physical_region_routes) + .await?; + } // Deletes metadata for table permanently. let table_route_value = TableRouteValue::new( self.table_id, @@ -157,7 +172,10 @@ mod tests { use crate::ddl::drop_database::cursor::DropDatabaseCursor; use crate::ddl::drop_database::executor::DropDatabaseExecutor; use crate::ddl::drop_database::{DropDatabaseContext, DropTableTarget, State}; - use crate::ddl::test_util::{create_logical_table, create_physical_table}; + use crate::ddl::test_util::datanode_handler::DatanodeWatcher; + use crate::ddl::test_util::{ + create_logical_table, create_physical_table, put_datanode_address, + }; use crate::error::{self, Error, Result}; use crate::key::datanode_table::DatanodeTableKey; use crate::peer::Peer; @@ -206,6 +224,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); assert!(!status.need_persist()); @@ -218,6 +237,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let mut state = DropDatabaseExecutor::new( physical_table_id, @@ -258,6 +278,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); assert!(!status.need_persist()); @@ -270,6 +291,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let mut state = DropDatabaseExecutor::new( logical_table_id, @@ -360,6 +382,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let err = state.next(&ddl_context, &mut ctx).await.unwrap_err(); assert!(err.is_retry_later()); @@ -389,6 +412,7 @@ mod tests { schema: DEFAULT_SCHEMA_NAME.to_string(), drop_if_exists: false, tables: None, + retrying: false, }; state.recover(&ddl_context).unwrap(); assert_eq!(state.dropping_regions.len(), 1); @@ -398,4 +422,45 @@ mod tests { assert_eq!(cursor.target, DropTableTarget::Physical); } } + + #[tokio::test] + async fn test_next_remaps_addresses_when_retrying() { + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let physical_table_id = create_physical_table(&ddl_context, "phy").await; + let (_, table_route) = ddl_context + .table_metadata_manager + .table_route_manager() + .get_physical_table_route(physical_table_id) + .await + .unwrap(); + + let mut state = DropDatabaseExecutor::new( + physical_table_id, + physical_table_id, + TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "phy"), + table_route.region_routes, + DropTableTarget::Physical, + ); + state.physical_region_routes[0] + .leader_peer + .as_mut() + .unwrap() + .addr = "old-addr".to_string(); + let mut ctx = DropDatabaseContext { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + drop_if_exists: false, + tables: None, + retrying: true, + }; + + put_datanode_address(&ddl_context, 0, "new-addr").await; + + state.next(&ddl_context, &mut ctx).await.unwrap(); + + let (peer, _) = rx.try_recv().unwrap(); + assert_eq!(peer.addr, "new-addr"); + } } diff --git a/src/common/meta/src/ddl/drop_database/metadata.rs b/src/common/meta/src/ddl/drop_database/metadata.rs index c71eded7af..0306dd2a3b 100644 --- a/src/common/meta/src/ddl/drop_database/metadata.rs +++ b/src/common/meta/src/ddl/drop_database/metadata.rs @@ -122,6 +122,7 @@ mod tests { schema: "bar".to_string(), drop_if_exists: true, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); state @@ -150,6 +151,7 @@ mod tests { schema: "bar".to_string(), drop_if_exists: true, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); state diff --git a/src/common/meta/src/ddl/drop_database/start.rs b/src/common/meta/src/ddl/drop_database/start.rs index 4da83e367f..6b3e149877 100644 --- a/src/common/meta/src/ddl/drop_database/start.rs +++ b/src/common/meta/src/ddl/drop_database/start.rs @@ -93,6 +93,7 @@ mod tests { schema: "bar".to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let err = step.next(&ddl_context, &mut ctx).await.unwrap_err(); assert_matches!(err, error::Error::SchemaNotFound { .. }); @@ -108,6 +109,7 @@ mod tests { schema: "bar".to_string(), drop_if_exists: true, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); state.as_any().downcast_ref::().unwrap(); @@ -130,6 +132,7 @@ mod tests { schema: "bar".to_string(), drop_if_exists: false, tables: None, + retrying: false, }; let (state, status) = state.next(&ddl_context, &mut ctx).await.unwrap(); state.as_any().downcast_ref::().unwrap(); diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 55c33330c4..8bd7c7155c 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -154,7 +154,24 @@ impl DropTableProcedure { Ok(Status::executing(true)) } - pub async fn on_datanode_drop_regions(&mut self) -> Result { + pub async fn on_datanode_drop_regions(&mut self, retrying: bool) -> Result { + if retrying { + info!( + "Remapping region routes addresses for retrying drop regions for table_id: {}", + self.data.table_id() + ); + let storage = self + .context + .table_metadata_manager + .table_route_manager() + .table_route_storage(); + // The peer addresses may change during retries, + // so we always remap the region routes. + storage + .remap_region_routes(&mut self.data.physical_region_routes) + .await?; + } + self.executor .on_drop_regions( &self.context.node_manager, @@ -215,7 +232,7 @@ impl Procedure for DropTableProcedure { Ok(()) } - async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult { + async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult { let state = &self.data.state; let _timer = metrics::METRIC_META_PROCEDURE_DROP_TABLE .with_label_values(&[state.as_ref()]) @@ -225,7 +242,10 @@ impl Procedure for DropTableProcedure { DropTableState::Prepare => self.on_prepare().await, DropTableState::DeleteMetadata => self.on_delete_metadata().await, DropTableState::InvalidateTableCache => self.on_broadcast().await, - DropTableState::DatanodeDropRegions => self.on_datanode_drop_regions().await, + DropTableState::DatanodeDropRegions => { + let retrying = ctx.is_retrying().await.unwrap_or(false); + self.on_datanode_drop_regions(retrying).await + } DropTableState::DeleteTombstone => self.on_delete_metadata_tombstone().await, } .map_err(map_to_procedure_error) diff --git a/src/common/meta/src/ddl/test_util.rs b/src/common/meta/src/ddl/test_util.rs index 36d422216f..7d8cdd895b 100644 --- a/src/common/meta/src/ddl/test_util.rs +++ b/src/common/meta/src/ddl/test_util.rs @@ -41,8 +41,12 @@ use crate::ddl::test_util::create_table::{ TestCreateTableExprBuilder, build_raw_table_info_from_expr, }; use crate::ddl::{DdlContext, TableMetadata}; +use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::table_route::TableRouteValue; +use crate::key::{MetadataKey, MetadataValue}; +use crate::peer::Peer; use crate::rpc::ddl::CreateTableTask; +use crate::rpc::store::PutRequest; pub async fn create_physical_table_metadata( ddl_context: &DdlContext, @@ -56,6 +60,21 @@ pub async fn create_physical_table_metadata( .unwrap(); } +pub async fn put_datanode_address(ddl_context: &DdlContext, node_id: u64, addr: &str) { + ddl_context + .table_metadata_manager + .kv_backend() + .put(PutRequest { + key: NodeAddressKey::with_datanode(node_id).to_bytes(), + value: NodeAddressValue::new(Peer::new(node_id, addr)) + .try_as_raw_value() + .unwrap(), + ..Default::default() + }) + .await + .unwrap(); +} + pub async fn create_physical_table(ddl_context: &DdlContext, name: &str) -> TableId { // Prepares physical table metadata. let mut create_physical_table_task = test_create_physical_table_task(name); diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index 5355ac8c7c..7f4a6bd716 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -42,7 +42,7 @@ use crate::ddl::test_util::datanode_handler::{ DatanodeWatcher, NaiveDatanodeHandler, RetryErrorDatanodeHandler, UnexpectedErrorDatanodeHandler, }; -use crate::ddl::test_util::{assert_column_name, get_raw_table_info}; +use crate::ddl::test_util::{assert_column_name, get_raw_table_info, put_datanode_address}; use crate::error::{Error, Result}; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; @@ -244,6 +244,27 @@ async fn test_on_datanode_create_regions_should_not_retry() { assert!(!error.is_retry_later()); } +#[tokio::test] +async fn test_on_datanode_create_regions_remaps_addresses_when_retrying() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx).with_handler(create_request_handler); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let ddl_context = new_ddl_context(node_manager); + let task = test_create_table_task("foo"); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); + procedure.on_prepare().await.unwrap(); + + let table_route = procedure.data.table_route.as_mut().unwrap(); + let leader = table_route.region_routes[0].leader_peer.as_mut().unwrap(); + leader.addr = "old-addr".to_string(); + put_datanode_address(&ddl_context, leader.id, "new-addr").await; + + procedure.on_datanode_create_regions(true).await.unwrap(); + + let (peer, _) = rx.try_recv().unwrap(); + assert_eq!(peer.addr, "new-addr"); +} + #[tokio::test] async fn test_on_create_metadata_error() { common_telemetry::init_default_ut_logging(); diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index fb2c882da0..ae81bd7f52 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -34,7 +34,7 @@ use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::test_util::datanode_handler::{DatanodeWatcher, NaiveDatanodeHandler}; use crate::ddl::test_util::{ create_logical_table, create_physical_table, create_physical_table_metadata, - test_create_logical_table_task, test_create_physical_table_task, + put_datanode_address, test_create_logical_table_task, test_create_physical_table_task, }; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; @@ -146,7 +146,7 @@ async fn test_on_datanode_drop_regions() { // Drop table let mut procedure = DropTableProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); - procedure.on_datanode_drop_regions().await.unwrap(); + procedure.on_datanode_drop_regions(false).await.unwrap(); let check = |peer: Peer, request: RegionRequest, @@ -186,6 +186,50 @@ async fn test_on_datanode_drop_regions() { check(peer, request, 5, RegionId::new(table_id, 1), true); } +#[tokio::test] +async fn test_on_datanode_drop_regions_remaps_addresses_when_retrying() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::new(1, "old-leader")), + follower_peers: vec![Peer::new(5, "old-follower")], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + + let task = new_drop_table_task(table_name, table_id, false); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); + procedure.on_prepare().await.unwrap(); + + put_datanode_address(&ddl_context, 1, "new-leader").await; + put_datanode_address(&ddl_context, 5, "new-follower").await; + + procedure.on_datanode_drop_regions(true).await.unwrap(); + + let mut peers = Vec::new(); + for _ in 0..2 { + peers.push(rx.try_recv().unwrap().0); + } + peers.sort_unstable_by_key(|p| p.id); + assert_eq!(peers[0].addr, "new-leader"); + assert_eq!(peers[1].addr, "new-follower"); +} + #[tokio::test] async fn test_on_rollback() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 40e8980a1c..21491b6230 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -42,7 +42,7 @@ use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; use crate::rpc::ddl::TruncateTableTask; -use crate::rpc::router::{RegionRoute, find_leader_regions, find_leaders}; +use crate::rpc::router::{find_leader_regions, find_leaders}; pub struct TruncateTableProcedure { context: DdlContext, @@ -94,12 +94,11 @@ impl TruncateTableProcedure { pub(crate) fn new( task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - region_routes: Vec, context: DdlContext, ) -> Self { Self { context, - data: TruncateTableData::new(task, table_info_value, region_routes), + data: TruncateTableData::new(task, table_info_value), } } @@ -138,13 +137,18 @@ impl TruncateTableProcedure { async fn on_datanode_truncate_regions(&mut self) -> Result { let table_id = self.data.table_id(); - let region_routes = &self.data.region_routes; - let leaders = find_leaders(region_routes); + let (_, physical_table_route) = self + .context + .table_metadata_manager + .table_route_manager() + .get_physical_table_route(table_id) + .await?; + let leaders = find_leaders(&physical_table_route.region_routes); let mut truncate_region_tasks = Vec::with_capacity(leaders.len()); for datanode in leaders { let requester = self.context.node_manager.datanode(&datanode).await; - let regions = find_leader_regions(region_routes, &datanode); + let regions = find_leader_regions(&physical_table_route.region_routes, &datanode); for region in regions { let region_id = RegionId::new(table_id, region); @@ -201,20 +205,17 @@ pub struct TruncateTableData { state: TruncateTableState, task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - region_routes: Vec, } impl TruncateTableData { pub fn new( task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - region_routes: Vec, ) -> Self { Self { state: TruncateTableState::Prepare, task, table_info_value, - region_routes, } } diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 0106add32f..d0619ca74f 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -45,7 +45,7 @@ use crate::ddl::drop_view::DropViewProcedure; use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::{DdlContext, utils}; use crate::error::{ - CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, ProcedureOutputSnafu, + self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, WaitProcedureSnafu, @@ -72,7 +72,6 @@ use crate::rpc::ddl::{ CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TruncateTableTask, }; -use crate::rpc::router::RegionRoute; /// A configurator that customizes or enhances a [`DdlManager`]. #[async_trait::async_trait] @@ -521,15 +520,9 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - region_routes: Vec, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = TruncateTableProcedure::new( - truncate_table_task, - table_info_value, - region_routes, - context, - ); + let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); @@ -658,19 +651,26 @@ async fn handle_truncate_table_task( let table_metadata_manager = &ddl_manager.table_metadata_manager(); let table_ref = truncate_table_task.table_ref(); - let (table_info_value, table_route_value) = - table_metadata_manager.get_full_table_info(table_id).await?; - - let table_info_value = table_info_value.with_context(|| TableInfoNotFoundSnafu { - table: table_ref.to_string(), - })?; - - let table_route_value = table_route_value.context(TableRouteNotFoundSnafu { table_id })?; - - let table_route = table_route_value.into_inner().region_routes()?.clone(); + let table_info_value = table_metadata_manager + .table_info_manager() + .get(table_id) + .await? + .with_context(|| TableInfoNotFoundSnafu { + table: table_ref.to_string(), + })?; + let physical_table_id = table_metadata_manager + .table_route_manager() + .get_physical_table_id(table_id) + .await?; + ensure!( + physical_table_id == table_id, + error::UnexpectedSnafu { + err_msg: "Truncate table is only supported for physical tables." + } + ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, table_route) + .submit_truncate_table_task(truncate_table_task, table_info_value) .await?; info!("Table: {table_id} is truncated via procedure_id {id:?}"); diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 332c60f225..3a9217e710 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -663,7 +663,7 @@ impl TableMetadataManager { if let Some(table_route_value) = &mut table_route_value { self.table_route_manager() .table_route_storage() - .remap_route_address(table_route_value) + .remap_table_route(table_route_value) .await?; } Ok((table_info_value, table_route_value)) diff --git a/src/common/meta/src/key/table_route.rs b/src/common/meta/src/key/table_route.rs index a409b75fce..4c618addda 100644 --- a/src/common/meta/src/key/table_route.rs +++ b/src/common/meta/src/key/table_route.rs @@ -675,7 +675,7 @@ impl TableRouteStorage { pub async fn get(&self, table_id: TableId) -> Result> { let mut table_route = self.get_inner(table_id).await?; if let Some(table_route) = &mut table_route { - self.remap_route_address(table_route).await?; + self.remap_table_route(table_route).await?; }; Ok(table_route) @@ -697,7 +697,7 @@ impl TableRouteStorage { ) -> Result>> { let mut table_route = self.get_with_raw_bytes_inner(table_id).await?; if let Some(table_route) = &mut table_route { - self.remap_route_address(table_route).await?; + self.remap_table_route(table_route).await?; }; Ok(table_route) @@ -791,10 +791,7 @@ impl TableRouteStorage { Ok(()) } - pub(crate) async fn remap_route_address( - &self, - table_route: &mut TableRouteValue, - ) -> Result<()> { + pub(crate) async fn remap_table_route(&self, table_route: &mut TableRouteValue) -> Result<()> { let keys = extract_address_keys(table_route).into_iter().collect(); let node_addrs = self.get_node_addresses(keys).await?; set_addresses(&node_addrs, table_route)?; @@ -802,6 +799,17 @@ impl TableRouteStorage { Ok(()) } + pub(crate) async fn remap_region_routes( + &self, + region_routes: &mut [RegionRoute], + ) -> Result<()> { + let keys = extract_address_keys_from_region_routes(region_routes) + .into_iter() + .collect(); + let node_addrs = self.get_node_addresses(keys).await?; + set_addresses_for_region_routes(&node_addrs, region_routes) + } + async fn get_node_addresses( &self, keys: Vec>, @@ -824,15 +832,11 @@ impl TableRouteStorage { } } -fn set_addresses( +fn set_addresses_for_region_routes( node_addrs: &HashMap, - table_route: &mut TableRouteValue, + region_routes: &mut [RegionRoute], ) -> Result<()> { - let TableRouteValue::Physical(physical_table_route) = table_route else { - return Ok(()); - }; - - for region_route in &mut physical_table_route.region_routes { + for region_route in region_routes { if let Some(leader) = &mut region_route.leader_peer && let Some(node_addr) = node_addrs.get(&leader.id) { @@ -848,13 +852,18 @@ fn set_addresses( Ok(()) } -fn extract_address_keys(table_route: &TableRouteValue) -> HashSet> { +fn set_addresses( + node_addrs: &HashMap, + table_route: &mut TableRouteValue, +) -> Result<()> { let TableRouteValue::Physical(physical_table_route) = table_route else { - return HashSet::default(); + return Ok(()); }; + set_addresses_for_region_routes(node_addrs, &mut physical_table_route.region_routes) +} - physical_table_route - .region_routes +fn extract_address_keys_from_region_routes(region_routes: &[RegionRoute]) -> HashSet> { + region_routes .iter() .flat_map(|region_route| { region_route @@ -871,6 +880,14 @@ fn extract_address_keys(table_route: &TableRouteValue) -> HashSet> { .collect() } +fn extract_address_keys(table_route: &TableRouteValue) -> HashSet> { + let TableRouteValue::Physical(physical_table_route) = table_route else { + return HashSet::default(); + }; + + extract_address_keys_from_region_routes(&physical_table_route.region_routes) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -1104,7 +1121,7 @@ mod tests { .unwrap(); table_route_storage - .remap_route_address(&mut table_route) + .remap_table_route(&mut table_route) .await .unwrap(); diff --git a/src/common/procedure/src/local/runner.rs b/src/common/procedure/src/local/runner.rs index 2a974de889..ca3e221f43 100644 --- a/src/common/procedure/src/local/runner.rs +++ b/src/common/procedure/src/local/runner.rs @@ -1293,6 +1293,60 @@ mod tests { .await; } + #[tokio::test] + async fn test_retrying_state_visible_in_context_on_retry() { + let retrying_states = Arc::new(std::sync::Mutex::new(Vec::new())); + let captured = retrying_states.clone(); + let mut times = 0; + + let exec_fn = move |ctx: Context| { + times += 1; + let captured = captured.clone(); + async move { + let is_retrying = ctx.is_retrying().await; + captured.lock().unwrap().push(is_retrying); + if times == 1 { + Err(Error::retry_later(MockError::new(StatusCode::Unexpected))) + } else { + Ok(Status::done()) + } + } + .boxed() + }; + + let procedure = ProcedureAdapter { + data: "retrying_state".to_string(), + lock_key: LockKey::single_exclusive("catalog.schema.table"), + poison_keys: PoisonKeys::default(), + exec_fn, + rollback_fn: None, + }; + + let dir = create_temp_dir("retrying_state"); + let meta = procedure.new_meta(ROOT_ID); + let object_store = test_util::new_object_store(&dir); + let procedure_store = Arc::new(ProcedureStore::from_object_store(object_store)); + let mut runner = new_runner(meta.clone(), Box::new(procedure), procedure_store); + let ctx = context_with_provider( + meta.id, + runner.manager_ctx.clone() as Arc, + ); + + runner + .manager_ctx + .procedures + .write() + .unwrap() + .insert(meta.id, runner.meta.clone()); + runner.manager_ctx.start(); + + runner.execute_once(&ctx).await; + runner.execute_once(&ctx).await; + + let states = retrying_states.lock().unwrap().clone(); + assert_eq!(states, vec![Some(false), Some(true)]); + } + #[tokio::test(flavor = "multi_thread")] async fn test_execute_on_retry_later_error_with_child() { common_telemetry::init_default_ut_logging(); diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index 843052ed97..8e34f4bb30 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -177,6 +177,18 @@ pub struct Context { pub provider: ContextProviderRef, } +impl Context { + /// Returns true if current procedure state is retrying. + pub async fn is_retrying(&self) -> Option { + self.provider + .procedure_state(self.procedure_id) + .await + .ok() + .flatten() + .map(|s| s.is_retrying()) + } +} + /// A `Procedure` represents an operation or a set of operations to be performed step-by-step. #[async_trait] pub trait Procedure: Send { diff --git a/src/meta-srv/src/procedure/tests.rs b/src/meta-srv/src/procedure/tests.rs index 105bfb7bc4..93879d3c08 100644 --- a/src/meta-srv/src/procedure/tests.rs +++ b/src/meta-srv/src/procedure/tests.rs @@ -219,7 +219,7 @@ async fn test_on_datanode_create_regions() { } }); - let status = procedure.on_datanode_create_regions().await.unwrap(); + let status = procedure.on_datanode_create_regions(false).await.unwrap(); assert!(matches!( status, Status::Executing { From 24ab861052f8baef6bb46f904e3bedeca4b201c2 Mon Sep 17 00:00:00 2001 From: Lanqing Yang Date: Thu, 9 Apr 2026 04:12:05 -0700 Subject: [PATCH 003/120] chore: move Tantivy fulltext search to blocking thread pool (#7919) perf: move Tantivy fulltext search to blocking thread pool Wrap the synchronous Tantivy search (query parsing, posting list traversal, stored field reads) in spawn_blocking_global to avoid starving the tokio async runtime with CPU-bound work. Signed-off-by: lyang24 --- .../src/fulltext_index/search/tantivy.rs | 123 ++++++++++-------- 1 file changed, 69 insertions(+), 54 deletions(-) diff --git a/src/index/src/fulltext_index/search/tantivy.rs b/src/index/src/fulltext_index/search/tantivy.rs index d06a5cd329..9fb0091868 100644 --- a/src/index/src/fulltext_index/search/tantivy.rs +++ b/src/index/src/fulltext_index/search/tantivy.rs @@ -14,6 +14,7 @@ use std::collections::{BTreeSet, HashMap}; use std::path::Path; +use std::sync::Arc; use std::time::Instant; use async_trait::async_trait; @@ -27,15 +28,19 @@ use tantivy::{Index, IndexReader, ReloadPolicy, TantivyDocument}; use crate::fulltext_index::Config; use crate::fulltext_index::create::{ROWID_FIELD_NAME, TEXT_FIELD_NAME}; use crate::fulltext_index::error::{ - Result, TantivyDocNotFoundSnafu, TantivyParserSnafu, TantivySnafu, + JoinSnafu, Result, TantivyDocNotFoundSnafu, TantivyParserSnafu, TantivySnafu, }; use crate::fulltext_index::search::{FulltextIndexSearcher, RowId}; /// `TantivyFulltextIndexSearcher` is a searcher using Tantivy. pub struct TantivyFulltextIndexSearcher { - /// Tanitvy index. + inner: Arc, +} + +struct TantivySearcherInner { + /// Tantivy index. index: Index, - /// Tanitvy index reader. + /// Tantivy index reader. reader: IndexReader, /// The default field used to build `QueryParser` default_field: Field, @@ -66,63 +71,73 @@ impl TantivyFulltextIndexSearcher { ); Ok(Self { - index, - reader, - default_field, + inner: Arc::new(TantivySearcherInner { + index, + reader, + default_field, + }), }) } } +fn search_sync(inner: &TantivySearcherInner, query: &str) -> Result> { + let searcher = inner.reader.searcher(); + let query_parser = QueryParser::for_index(&inner.index, vec![inner.default_field]); + let query = query_parser + .parse_query(query) + .context(TantivyParserSnafu)?; + let doc_addrs = searcher + .search(&query, &DocSetCollector) + .context(TantivySnafu)?; + + let seg_metas = inner + .index + .searchable_segment_metas() + .context(TantivySnafu)?; + + // FAST PATH: only one segment, the doc id is the same as the row id. + // Also for compatibility with the old version. + if seg_metas.len() == 1 { + return Ok(doc_addrs.into_iter().map(|d| d.doc_id).collect()); + } + + // SLOW PATH: multiple segments, need to calculate the row id. + let rowid_field = searcher + .schema() + .get_field(ROWID_FIELD_NAME) + .context(TantivySnafu)?; + let mut seg_offsets = HashMap::with_capacity(seg_metas.len()); + let mut res = BTreeSet::new(); + for doc_addr in doc_addrs { + let offset = if let Some(offset) = seg_offsets.get(&doc_addr.segment_ord) { + *offset + } else { + // Calculate the offset at the first time meeting the segment and cache it since + // the offset is the same for all rows in the same segment. + let doc: TantivyDocument = searcher.doc(doc_addr).context(TantivySnafu)?; + let rowid = doc + .get_first(rowid_field) + .and_then(|v| v.as_u64()) + .context(TantivyDocNotFoundSnafu { doc_addr })?; + + let offset = rowid as u32 - doc_addr.doc_id; + seg_offsets.insert(doc_addr.segment_ord, offset); + offset + }; + + res.insert(doc_addr.doc_id + offset); + } + + Ok(res) +} + #[async_trait] impl FulltextIndexSearcher for TantivyFulltextIndexSearcher { async fn search(&self, query: &str) -> Result> { - let searcher = self.reader.searcher(); - let query_parser = QueryParser::for_index(&self.index, vec![self.default_field]); - let query = query_parser - .parse_query(query) - .context(TantivyParserSnafu)?; - let doc_addrs = searcher - .search(&query, &DocSetCollector) - .context(TantivySnafu)?; - - let seg_metas = self - .index - .searchable_segment_metas() - .context(TantivySnafu)?; - - // FAST PATH: only one segment, the doc id is the same as the row id. - // Also for compatibility with the old version. - if seg_metas.len() == 1 { - return Ok(doc_addrs.into_iter().map(|d| d.doc_id).collect()); - } - - // SLOW PATH: multiple segments, need to calculate the row id. - let rowid_field = searcher - .schema() - .get_field(ROWID_FIELD_NAME) - .context(TantivySnafu)?; - let mut seg_offsets = HashMap::with_capacity(seg_metas.len()); - let mut res = BTreeSet::new(); - for doc_addr in doc_addrs { - let offset = if let Some(offset) = seg_offsets.get(&doc_addr.segment_ord) { - *offset - } else { - // Calculate the offset at the first time meeting the segment and cache it since - // the offset is the same for all rows in the same segment. - let doc: TantivyDocument = searcher.doc(doc_addr).context(TantivySnafu)?; - let rowid = doc - .get_first(rowid_field) - .and_then(|v| v.as_u64()) - .context(TantivyDocNotFoundSnafu { doc_addr })?; - - let offset = rowid as u32 - doc_addr.doc_id; - seg_offsets.insert(doc_addr.segment_ord, offset); - offset - }; - - res.insert(doc_addr.doc_id + offset); - } - - Ok(res) + let inner = self.inner.clone(); + let query = query.to_string(); + common_runtime::spawn_blocking_global(move || search_sync(&inner, &query)) + .await + .context(JoinSnafu)? } } From fb5333e116d89fd073d2c40e84486caa7a5a83db Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 9 Apr 2026 20:37:24 +0800 Subject: [PATCH 004/120] ci: add standalone workflows for bumping helm charts and homebrew (#7941) ci: add standalone workflows for bumping helm charts and homebrew versions Signed-off-by: evenyag --- .../workflows/bump-helm-charts-version.yml | 29 +++++++++++++++++++ .../bump-homebrew-greptime-version.yml | 29 +++++++++++++++++++ 2 files changed, 58 insertions(+) create mode 100644 .github/workflows/bump-helm-charts-version.yml create mode 100644 .github/workflows/bump-homebrew-greptime-version.yml diff --git a/.github/workflows/bump-helm-charts-version.yml b/.github/workflows/bump-helm-charts-version.yml new file mode 100644 index 0000000000..5921ec8a8c --- /dev/null +++ b/.github/workflows/bump-helm-charts-version.yml @@ -0,0 +1,29 @@ +name: Bump helm charts version + +on: + workflow_dispatch: + inputs: + version: + description: The version to bump (e.g. v1.0.0) + required: true + type: string + +jobs: + bump-helm-charts-version: + name: Bump helm charts version + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Bump helm charts version + env: + GITHUB_TOKEN: ${{ secrets.HELM_CHARTS_REPO_TOKEN }} + VERSION: ${{ inputs.version }} + run: | + ./.github/scripts/update-helm-charts-version.sh diff --git a/.github/workflows/bump-homebrew-greptime-version.yml b/.github/workflows/bump-homebrew-greptime-version.yml new file mode 100644 index 0000000000..af8ca8fc99 --- /dev/null +++ b/.github/workflows/bump-homebrew-greptime-version.yml @@ -0,0 +1,29 @@ +name: Bump homebrew greptime version + +on: + workflow_dispatch: + inputs: + version: + description: The version to bump (e.g. v1.0.0) + required: true + type: string + +jobs: + bump-homebrew-greptime-version: + name: Bump homebrew greptime version + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Bump homebrew greptime version + env: + GITHUB_TOKEN: ${{ secrets.HOMEBREW_GREPTIME_REPO_TOKEN }} + VERSION: ${{ inputs.version }} + run: | + ./.github/scripts/update-homebrew-greptme-version.sh From e9d783cccf6b41596b4537ac102b9e757a462da8 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Fri, 10 Apr 2026 03:18:56 +0800 Subject: [PATCH 005/120] feat: execution timeout for prepared statement (#7932) * feat: execution timeout for prepared statement * fix: lint fix --- src/frontend/src/instance.rs | 58 +++++++++++++++++++++++++++++++----- 1 file changed, 50 insertions(+), 8 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index ce589bb677..99444bb2a2 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -483,6 +483,27 @@ fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option, + query_ctx: &QueryContextRef, +) -> Option { + match stmt { + Some(s) => derive_timeout(s, query_ctx), + None => { + let query_timeout = query_ctx.query_timeout()?; + if query_timeout.is_zero() { + return None; + } + match query_ctx.channel() { + Channel::Postgres => Some(query_timeout), + _ => None, + } + } + } +} + fn attach_timeout(output: Output, mut timeout: Duration) -> Result { if timeout.is_zero() { return StatementTimeoutSnafu.fail(); @@ -588,6 +609,33 @@ impl Instance { } } + async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { + self.query_engine + .execute(plan, query_ctx) + .await + .context(ExecLogicalPlanSnafu) + } + + async fn exec_plan_with_timeout( + &self, + stmt: Option, + plan: LogicalPlan, + query_ctx: QueryContextRef, + ) -> Result { + let timeout = derive_timeout_for_plan(stmt.as_ref(), &query_ctx); + match timeout { + Some(timeout) => { + let start = tokio::time::Instant::now(); + let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx)) + .await + .map_err(|_| StatementTimeoutSnafu.build())??; + let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default(); + attach_timeout(output, remaining_timeout) + } + None => self.exec_plan(plan, query_ctx).await, + } + } + async fn do_exec_plan_inner( &self, stmt: Option, @@ -624,7 +672,7 @@ impl Instance { slow_query_timer, ); - let query_fut = self.query_engine.execute(plan.clone(), query_ctx); + let query_fut = self.exec_plan_with_timeout(Some(stmt), plan, query_ctx); CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) .await @@ -640,14 +688,8 @@ impl Instance { }; Output { data, meta } }) - .context(ExecLogicalPlanSnafu) } else { - // plan should be prepared before exec - // we'll do check there - self.query_engine - .execute(plan.clone(), query_ctx) - .await - .context(ExecLogicalPlanSnafu) + self.exec_plan_with_timeout(stmt, plan, query_ctx).await } } From fd94f5519398af27f9fb7119abffa43ac0315e70 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 10 Apr 2026 11:12:33 +0800 Subject: [PATCH 006/120] refactor(mito2): remove dead scan code (#7925) * refactor(mito2): remove dead batch parallel scan helpers Signed-off-by: evenyag * refactor(mito2): remove dead merge reader path Signed-off-by: evenyag * refactor(mito2): remove dead batch dedup reader Signed-off-by: evenyag * test(mito2): remove obsolete batch source helper Signed-off-by: evenyag * refactor: remove unused plain batch Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/benches/simple_bulk_memtable.rs | 47 +- .../src/memtable/simple_bulk_memtable.rs | 106 +- src/mito2/src/read.rs | 2 - src/mito2/src/read/dedup.rs | 544 +--------- src/mito2/src/read/flat_merge.rs | 82 +- src/mito2/src/read/merge.rs | 982 ------------------ src/mito2/src/read/plain_batch.rs | 505 --------- src/mito2/src/read/scan_region.rs | 78 +- src/mito2/src/read/scan_util.rs | 2 +- src/mito2/src/sst.rs | 28 - src/mito2/src/test_util/sst_util.rs | 10 +- 11 files changed, 94 insertions(+), 2292 deletions(-) delete mode 100644 src/mito2/src/read/merge.rs delete mode 100644 src/mito2/src/read/plain_batch.rs diff --git a/src/mito2/benches/simple_bulk_memtable.rs b/src/mito2/benches/simple_bulk_memtable.rs index 05035734de..8a199f46f1 100644 --- a/src/mito2/benches/simple_bulk_memtable.rs +++ b/src/mito2/benches/simple_bulk_memtable.rs @@ -21,11 +21,7 @@ use criterion::{Criterion, criterion_group, criterion_main}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; use mito2::memtable::simple_bulk_memtable::SimpleBulkMemtable; -use mito2::memtable::{IterBuilder, KeyValues, Memtable, MemtableRanges, RangesOptions}; -use mito2::read; -use mito2::read::Source; -use mito2::read::dedup::DedupReader; -use mito2::read::merge::MergeReaderBuilder; +use mito2::memtable::{IterBuilder, KeyValues, Memtable, RangesOptions}; use mito2::region::options::MergeMode; use mito2::test_util::column_metadata_to_column_schema; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; @@ -126,36 +122,6 @@ fn create_memtable_with_rows(num_batches: usize) -> SimpleBulkMemtable { } async fn flush(mem: &SimpleBulkMemtable) { - let MemtableRanges { ranges, .. } = mem.ranges(None, RangesOptions::for_flush()).unwrap(); - - let mut source = if ranges.len() == 1 { - let only_range = ranges.into_values().next().unwrap(); - let iter = only_range.build_iter().unwrap(); - Source::Iter(iter) - } else { - let sources = ranges - .into_values() - .map(|r| r.build_iter().map(Source::Iter)) - .collect::>>() - .unwrap(); - let merge_reader = MergeReaderBuilder::from_sources(sources) - .build() - .await - .unwrap(); - let reader = Box::new(DedupReader::new( - merge_reader, - read::dedup::LastRow::new(true), - None, - )); - Source::Reader(reader) - }; - - while let Some(b) = source.next_batch().await.unwrap() { - black_box(b); - } -} - -async fn flush_original(mem: &SimpleBulkMemtable) { let iter = mem .ranges(None, RangesOptions::default()) .unwrap() @@ -179,19 +145,10 @@ fn bench_ranges_parallel_vs_sequential(c: &mut Criterion) { let total_rows_k = num_batch * 10; let memtable = create_memtable_with_rows(num_batch); - group.bench_with_input( - BenchmarkId::new("flush_by_merge_reader", format!("{}k_rows", total_rows_k)), - &memtable, - |b, memtable| b.to_async(&rt).iter(|| async { flush(memtable).await }), - ); - group.bench_with_input( BenchmarkId::new("flush_by_iter", format!("{}k_rows", total_rows_k)), &memtable, - |b, memtable| { - b.to_async(&rt) - .iter(|| async { flush_original(memtable).await }) - }, + |b, memtable| b.to_async(&rt).iter(|| async { flush(memtable).await }), ); } diff --git a/src/mito2/src/memtable/simple_bulk_memtable.rs b/src/mito2/src/memtable/simple_bulk_memtable.rs index 1284741347..6ff799ebf5 100644 --- a/src/mito2/src/memtable/simple_bulk_memtable.rs +++ b/src/mito2/src/memtable/simple_bulk_memtable.rs @@ -421,10 +421,6 @@ mod tests { use store_api::storage::{RegionId, SequenceNumber, SequenceRange}; use super::*; - use crate::read; - use crate::read::dedup::DedupReader; - use crate::read::merge::MergeReaderBuilder; - use crate::read::{BatchReader, Source}; use crate::region::options::MergeMode; use crate::test_util::column_metadata_to_column_schema; @@ -621,81 +617,6 @@ mod tests { assert_eq!(1, batch.num_rows()); } - #[tokio::test] - async fn test_write_dedup() { - let memtable = new_test_memtable(true, MergeMode::LastRow); - let kvs = build_key_values( - &memtable.region_metadata, - 0, - &[(1, 1.0, "a".to_string())], - OpType::Put, - ); - let kv = kvs.iter().next().unwrap(); - memtable.write_one(kv).unwrap(); - memtable.freeze().unwrap(); - - let kvs = build_key_values( - &memtable.region_metadata, - 1, - &[(1, 1.0, "a".to_string())], - OpType::Delete, - ); - let kv = kvs.iter().next().unwrap(); - memtable.write_one(kv).unwrap(); - - let ranges = memtable.ranges(None, RangesOptions::default()).unwrap(); - let mut source = vec![]; - for r in ranges.ranges.values() { - source.push(Source::Iter(r.build_iter().unwrap())); - } - - let reader = MergeReaderBuilder::from_sources(source) - .build() - .await - .unwrap(); - - let mut reader = DedupReader::new(reader, read::dedup::LastRow::new(false), None); - let mut num_rows = 0; - while let Some(b) = reader.next_batch().await.unwrap() { - num_rows += b.num_rows(); - } - assert_eq!(num_rows, 1); - } - - #[tokio::test] - async fn test_delete_only() { - let memtable = new_test_memtable(true, MergeMode::LastRow); - let kvs = build_key_values( - &memtable.region_metadata, - 0, - &[(1, 1.0, "a".to_string())], - OpType::Delete, - ); - let kv = kvs.iter().next().unwrap(); - memtable.write_one(kv).unwrap(); - memtable.freeze().unwrap(); - - let ranges = memtable.ranges(None, RangesOptions::default()).unwrap(); - let mut source = vec![]; - for r in ranges.ranges.values() { - source.push(Source::Iter(r.build_iter().unwrap())); - } - - let reader = MergeReaderBuilder::from_sources(source) - .build() - .await - .unwrap(); - - let mut reader = DedupReader::new(reader, read::dedup::LastRow::new(false), None); - let mut num_rows = 0; - while let Some(b) = reader.next_batch().await.unwrap() { - num_rows += b.num_rows(); - assert_eq!(b.num_rows(), 1); - assert_eq!(b.op_types().get_data(0).unwrap(), OpType::Delete as u8); - } - assert_eq!(num_rows, 1); - } - #[tokio::test] async fn test_single_range() { let memtable = new_test_memtable(true, MergeMode::LastRow); @@ -902,8 +823,8 @@ mod tests { .unwrap() } - #[tokio::test] - async fn test_write_read_large_string() { + #[test] + fn test_write_read_large_string() { let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); builder .push_column_metadata(ColumnMetadata { @@ -948,25 +869,12 @@ mod tests { .unwrap(); let MemtableRanges { ranges, .. } = memtable.ranges(None, RangesOptions::default()).unwrap(); - let mut source = if ranges.len() == 1 { - let only_range = ranges.into_values().next().unwrap(); - Source::Iter(only_range.build_iter().unwrap()) - } else { - let sources = ranges - .into_values() - .map(|r| r.build_iter().map(Source::Iter)) - .collect::>>() - .unwrap(); - let merge_reader = MergeReaderBuilder::from_sources(sources) - .build() - .await - .unwrap(); - Source::Reader(Box::new(merge_reader)) - }; - let mut rows = 0; - while let Some(b) = source.next_batch().await.unwrap() { - rows += b.num_rows(); + for range in ranges.into_values() { + let iter = range.build_iter().unwrap(); + for batch in iter { + rows += batch.unwrap().num_rows(); + } } assert_eq!(rows, 2); } diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index db7dfd1958..90eb9a3da7 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -21,8 +21,6 @@ pub mod flat_dedup; pub mod flat_merge; pub mod flat_projection; pub mod last_row; -pub mod merge; -pub mod plain_batch; pub mod projection; pub(crate) mod prune; pub(crate) mod pruner; diff --git a/src/mito2/src/read/dedup.rs b/src/mito2/src/read/dedup.rs index 5c881459b2..86f6b07ffc 100644 --- a/src/mito2/src/read/dedup.rs +++ b/src/mito2/src/read/dedup.rs @@ -19,17 +19,13 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::OpType; -use async_trait::async_trait; -use common_telemetry::debug; -use common_time::Timestamp; use datatypes::data_type::DataType; use datatypes::prelude::ScalarVector; use datatypes::value::Value; use datatypes::vectors::MutableVector; use crate::error::Result; -use crate::metrics::MERGE_FILTER_ROWS_TOTAL; -use crate::read::{Batch, BatchColumn, BatchReader}; +use crate::read::{Batch, BatchColumn}; /// Trait for reporting dedup metrics. pub trait DedupMetricsReport: Send + Sync { @@ -37,80 +33,6 @@ pub trait DedupMetricsReport: Send + Sync { fn report(&self, metrics: &mut DedupMetrics); } -/// A reader that dedup sorted batches from a source based on the -/// dedup strategy. -pub struct DedupReader { - source: R, - strategy: S, - metrics: DedupMetrics, - /// Optional metrics reporter. - metrics_reporter: Option>, -} - -impl DedupReader { - /// Creates a new dedup reader. - pub fn new( - source: R, - strategy: S, - metrics_reporter: Option>, - ) -> Self { - Self { - source, - strategy, - metrics: DedupMetrics::default(), - metrics_reporter, - } - } -} - -impl DedupReader { - /// Returns the next deduplicated batch. - async fn fetch_next_batch(&mut self) -> Result> { - while let Some(batch) = self.source.next_batch().await? { - if let Some(batch) = self.strategy.push_batch(batch, &mut self.metrics)? { - self.metrics.maybe_report(&self.metrics_reporter); - return Ok(Some(batch)); - } - } - - let result = self.strategy.finish(&mut self.metrics)?; - self.metrics.maybe_report(&self.metrics_reporter); - Ok(result) - } -} - -#[async_trait] -impl BatchReader for DedupReader { - async fn next_batch(&mut self) -> Result> { - self.fetch_next_batch().await - } -} - -impl Drop for DedupReader { - fn drop(&mut self) { - debug!("Dedup reader finished, metrics: {:?}", self.metrics); - - MERGE_FILTER_ROWS_TOTAL - .with_label_values(&["dedup"]) - .inc_by(self.metrics.num_unselected_rows as u64); - MERGE_FILTER_ROWS_TOTAL - .with_label_values(&["delete"]) - .inc_by(self.metrics.num_unselected_rows as u64); - - // Report any remaining metrics. - if let Some(reporter) = &self.metrics_reporter { - reporter.report(&mut self.metrics); - } - } -} - -#[cfg(test)] -impl DedupReader { - fn metrics(&self) -> &DedupMetrics { - &self.metrics - } -} - /// Strategy to remove duplicate rows from sorted batches. pub trait DedupStrategy: Send { /// Pushes a batch to the dedup strategy. @@ -124,114 +46,6 @@ pub trait DedupStrategy: Send { fn finish(&mut self, metrics: &mut DedupMetrics) -> Result>; } -/// State of the last row in a batch for dedup. -struct BatchLastRow { - primary_key: Vec, - /// The last timestamp of the batch. - timestamp: Timestamp, -} - -/// Dedup strategy that keeps the row with latest sequence of each key. -/// -/// This strategy is optimized specially based on the properties of the SST files, -/// memtables and the merge reader. It assumes that batches from files and memtables -/// don't contain duplicate rows and the merge reader never concatenates batches from -/// different source. -/// -/// We might implement a new strategy if we need to process files with duplicate rows. -pub struct LastRow { - /// Meta of the last row in the previous batch that has the same key - /// as the batch to push. - prev_batch: Option, - /// Filter deleted rows. - filter_deleted: bool, -} - -impl LastRow { - /// Creates a new strategy with the given `filter_deleted` flag. - pub fn new(filter_deleted: bool) -> Self { - Self { - prev_batch: None, - filter_deleted, - } - } -} - -impl DedupStrategy for LastRow { - fn push_batch( - &mut self, - mut batch: Batch, - metrics: &mut DedupMetrics, - ) -> Result> { - let start = Instant::now(); - - if batch.is_empty() { - return Ok(None); - } - debug_assert!(batch.first_timestamp().is_some()); - let prev_timestamp = match &self.prev_batch { - Some(prev_batch) => { - if prev_batch.primary_key != batch.primary_key() { - // The key has changed. This is the first batch of the - // new key. - None - } else { - Some(prev_batch.timestamp) - } - } - None => None, - }; - if batch.first_timestamp() == prev_timestamp { - metrics.num_unselected_rows += 1; - // This batch contains a duplicate row, skip it. - if batch.num_rows() == 1 { - // We don't need to update `prev_batch` because they have the same - // key and timestamp. - metrics.dedup_cost += start.elapsed(); - return Ok(None); - } - // Skips the first row. - batch = batch.slice(1, batch.num_rows() - 1); - } - - // Store current batch to `prev_batch` so we could compare the next batch - // with this batch. We store batch before filtering it as rows with `OpType::Delete` - // would be removed from the batch after filter, then we may store an incorrect `last row` - // of previous batch. - match &mut self.prev_batch { - Some(prev) => { - // Reuse the primary key buffer. - prev.primary_key.clone_from(&batch.primary_key); - prev.timestamp = batch.last_timestamp().unwrap(); - } - None => { - self.prev_batch = Some(BatchLastRow { - primary_key: batch.primary_key().to_vec(), - timestamp: batch.last_timestamp().unwrap(), - }) - } - } - - // Filters deleted rows. - if self.filter_deleted { - filter_deleted_from_batch(&mut batch, metrics)?; - } - - metrics.dedup_cost += start.elapsed(); - - // The batch can become empty if all rows are deleted. - if batch.is_empty() { - Ok(None) - } else { - Ok(Some(batch)) - } - } - - fn finish(&mut self, _metrics: &mut DedupMetrics) -> Result> { - Ok(None) - } -} - /// Removes deleted rows from the batch and updates metrics. fn filter_deleted_from_batch(batch: &mut Batch, metrics: &mut DedupMetrics) -> Result<()> { let num_rows = batch.num_rows(); @@ -672,137 +486,10 @@ impl>> Iterator for LastNonNullIter { mod tests { use std::sync::Arc; - use api::v1::OpType; use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array}; use super::*; use crate::read::BatchBuilder; - use crate::test_util::{VecBatchReader, check_reader_result, new_batch}; - - #[tokio::test] - async fn test_dedup_reader_no_duplications() { - let input = [ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch(b"k1", &[3], &[13], &[OpType::Put], &[23]), - new_batch( - b"k2", - &[1, 2], - &[111, 112], - &[OpType::Put, OpType::Put], - &[31, 32], - ), - ]; - - // Test last row. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastRow::new(true), None); - check_reader_result(&mut reader, &input).await; - assert_eq!(0, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - - // Test last non-null. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(true), None); - check_reader_result(&mut reader, &input).await; - assert_eq!(0, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } - - #[tokio::test] - async fn test_dedup_reader_duplications() { - let input = [ - new_batch( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[11, 12], - ), - // empty batch. - new_batch(b"k1", &[], &[], &[], &[]), - // Duplicate with the previous batch. - new_batch( - b"k1", - &[2, 3, 4], - &[10, 13, 13], - &[OpType::Put, OpType::Put, OpType::Delete], - &[2, 13, 14], - ), - new_batch( - b"k2", - &[1, 2], - &[20, 20], - &[OpType::Put, OpType::Delete], - &[101, 0], - ), - new_batch(b"k2", &[2], &[19], &[OpType::Put], &[102]), - new_batch(b"k3", &[2], &[20], &[OpType::Put], &[202]), - // This batch won't increase the deleted rows count as it - // is filtered out by the previous batch. - new_batch(b"k3", &[2], &[19], &[OpType::Delete], &[0]), - ]; - // Filter deleted. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastRow::new(true), None); - check_reader_result( - &mut reader, - &[ - new_batch( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[11, 12], - ), - new_batch(b"k1", &[3], &[13], &[OpType::Put], &[13]), - new_batch(b"k2", &[1], &[20], &[OpType::Put], &[101]), - new_batch(b"k3", &[2], &[20], &[OpType::Put], &[202]), - ], - ) - .await; - assert_eq!(5, reader.metrics().num_unselected_rows); - assert_eq!(2, reader.metrics().num_deleted_rows); - - // Does not filter deleted. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastRow::new(false), None); - check_reader_result( - &mut reader, - &[ - new_batch( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[11, 12], - ), - new_batch( - b"k1", - &[3, 4], - &[13, 13], - &[OpType::Put, OpType::Delete], - &[13, 14], - ), - new_batch( - b"k2", - &[1, 2], - &[20, 20], - &[OpType::Put, OpType::Delete], - &[101, 0], - ), - new_batch(b"k3", &[2], &[20], &[OpType::Put], &[202]), - ], - ) - .await; - assert_eq!(3, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } /// Returns a new [Batch] whose field has column id 1, 2. fn new_batch_multi_fields( @@ -839,235 +526,6 @@ mod tests { builder.build().unwrap() } - #[tokio::test] - async fn test_last_non_null_merge() { - let input = [ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (None, None)], - ), - // empty batch. - new_batch_multi_fields(b"k1", &[], &[], &[], &[]), - // Duplicate with the previous batch. - new_batch_multi_fields(b"k1", &[2], &[10], &[OpType::Put], &[(Some(12), None)]), - new_batch_multi_fields( - b"k1", - &[2, 3, 4], - &[10, 13, 13], - &[OpType::Put, OpType::Put, OpType::Delete], - &[(Some(2), Some(22)), (Some(13), None), (None, Some(14))], - ), - new_batch_multi_fields( - b"k2", - &[1, 2], - &[20, 20], - &[OpType::Put, OpType::Delete], - &[(Some(101), Some(101)), (None, None)], - ), - new_batch_multi_fields( - b"k2", - &[2], - &[19], - &[OpType::Put], - &[(Some(102), Some(102))], - ), - new_batch_multi_fields( - b"k3", - &[2], - &[20], - &[OpType::Put], - &[(Some(202), Some(202))], - ), - // This batch won't increase the deleted rows count as it - // is filtered out by the previous batch. (All fields are null). - new_batch_multi_fields(b"k3", &[2], &[19], &[OpType::Delete], &[(None, None)]), - ]; - - // Filter deleted. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(true), None); - check_reader_result( - &mut reader, - &[ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (Some(12), Some(22))], - ), - new_batch_multi_fields(b"k1", &[3], &[13], &[OpType::Put], &[(Some(13), None)]), - new_batch_multi_fields( - b"k2", - &[1], - &[20], - &[OpType::Put], - &[(Some(101), Some(101))], - ), - new_batch_multi_fields( - b"k3", - &[2], - &[20], - &[OpType::Put], - &[(Some(202), Some(202))], - ), - ], - ) - .await; - assert_eq!(6, reader.metrics().num_unselected_rows); - assert_eq!(2, reader.metrics().num_deleted_rows); - - // Does not filter deleted. - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(false), None); - check_reader_result( - &mut reader, - &[ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (Some(12), Some(22))], - ), - new_batch_multi_fields( - b"k1", - &[3, 4], - &[13, 13], - &[OpType::Put, OpType::Delete], - &[(Some(13), None), (None, Some(14))], - ), - new_batch_multi_fields( - b"k2", - &[1, 2], - &[20, 20], - &[OpType::Put, OpType::Delete], - &[(Some(101), Some(101)), (None, None)], - ), - new_batch_multi_fields( - b"k3", - &[2], - &[20], - &[OpType::Put], - &[(Some(202), Some(202))], - ), - ], - ) - .await; - assert_eq!(4, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } - - #[tokio::test] - async fn test_last_non_null_skip_merge_single() { - let input = [new_batch_multi_fields( - b"k1", - &[1, 2, 3], - &[13, 11, 13], - &[OpType::Put, OpType::Delete, OpType::Put], - &[(Some(11), Some(11)), (None, None), (Some(13), Some(13))], - )]; - - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(true), None); - check_reader_result( - &mut reader, - &[new_batch_multi_fields( - b"k1", - &[1, 3], - &[13, 13], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (Some(13), Some(13))], - )], - ) - .await; - assert_eq!(1, reader.metrics().num_unselected_rows); - assert_eq!(1, reader.metrics().num_deleted_rows); - - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(false), None); - check_reader_result(&mut reader, &input).await; - assert_eq!(0, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } - - #[tokio::test] - async fn test_last_non_null_skip_merge_no_null() { - let input = [ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (Some(12), Some(12))], - ), - new_batch_multi_fields(b"k1", &[2], &[10], &[OpType::Put], &[(None, Some(22))]), - new_batch_multi_fields( - b"k1", - &[2, 3], - &[9, 13], - &[OpType::Put, OpType::Put], - &[(Some(32), None), (Some(13), Some(13))], - ), - ]; - - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(true), None); - check_reader_result( - &mut reader, - &[ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (Some(12), Some(12))], - ), - new_batch_multi_fields(b"k1", &[3], &[13], &[OpType::Put], &[(Some(13), Some(13))]), - ], - ) - .await; - assert_eq!(2, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } - - #[tokio::test] - async fn test_last_non_null_merge_null() { - let input = [ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (None, None)], - ), - new_batch_multi_fields(b"k1", &[2], &[10], &[OpType::Put], &[(None, Some(22))]), - new_batch_multi_fields(b"k1", &[3], &[13], &[OpType::Put], &[(Some(33), None)]), - ]; - - let reader = VecBatchReader::new(&input); - let mut reader = DedupReader::new(reader, LastNonNull::new(true), None); - check_reader_result( - &mut reader, - &[ - new_batch_multi_fields( - b"k1", - &[1, 2], - &[13, 11], - &[OpType::Put, OpType::Put], - &[(Some(11), Some(11)), (None, Some(22))], - ), - new_batch_multi_fields(b"k1", &[3], &[13], &[OpType::Put], &[(Some(33), None)]), - ], - ) - .await; - assert_eq!(1, reader.metrics().num_unselected_rows); - assert_eq!(0, reader.metrics().num_deleted_rows); - } - fn check_dedup_strategy(input: &[Batch], strategy: &mut dyn DedupStrategy, expect: &[Batch]) { let mut actual = Vec::new(); let mut metrics = DedupMetrics::default(); diff --git a/src/mito2/src/read/flat_merge.rs b/src/mito2/src/read/flat_merge.rs index 946f2a610c..b1c304f244 100644 --- a/src/mito2/src/read/flat_merge.rs +++ b/src/mito2/src/read/flat_merge.rs @@ -14,8 +14,9 @@ use std::cmp::Ordering; use std::collections::BinaryHeap; +use std::fmt; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use async_stream::try_stream; use common_telemetry::debug; @@ -34,7 +35,6 @@ use crate::error::{ComputeArrowSnafu, Result}; use crate::memtable::BoxedRecordBatchIterator; use crate::metrics::READ_STAGE_ELAPSED; use crate::read::BoxedRecordBatchStream; -use crate::read::merge::{MergeMetrics, MergeMetricsReport}; use crate::sst::parquet::flat_format::{ primary_key_column_index, sequence_column_index, time_index_column_index, }; @@ -105,6 +105,84 @@ struct BatchCursor { row_idx: usize, } +/// Trait for reporting merge metrics. +pub trait MergeMetricsReport: Send + Sync { + /// Reports and resets the metrics. + fn report(&self, metrics: &mut MergeMetrics); +} + +/// Metrics for the merge reader. +#[derive(Default)] +pub struct MergeMetrics { + /// Cost to initialize the reader. + pub(crate) init_cost: Duration, + /// Total scan cost of the reader. + pub(crate) scan_cost: Duration, + /// Number of times to fetch batches. + pub(crate) num_fetch_by_batches: usize, + /// Number of times to fetch rows. + pub(crate) num_fetch_by_rows: usize, + /// Cost to fetch batches from sources. + pub(crate) fetch_cost: Duration, +} + +impl fmt::Debug for MergeMetrics { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.scan_cost.is_zero() { + return write!(f, "{{}}"); + } + + write!(f, r#"{{"scan_cost":"{:?}""#, self.scan_cost)?; + + if !self.init_cost.is_zero() { + write!(f, r#", "init_cost":"{:?}""#, self.init_cost)?; + } + if self.num_fetch_by_batches > 0 { + write!( + f, + r#", "num_fetch_by_batches":{}"#, + self.num_fetch_by_batches + )?; + } + if self.num_fetch_by_rows > 0 { + write!(f, r#", "num_fetch_by_rows":{}"#, self.num_fetch_by_rows)?; + } + if !self.fetch_cost.is_zero() { + write!(f, r#", "fetch_cost":"{:?}""#, self.fetch_cost)?; + } + + write!(f, "}}") + } +} + +impl MergeMetrics { + /// Merges metrics from another MergeMetrics instance. + pub(crate) fn merge(&mut self, other: &MergeMetrics) { + let MergeMetrics { + init_cost, + scan_cost, + num_fetch_by_batches, + num_fetch_by_rows, + fetch_cost, + } = other; + + self.init_cost += *init_cost; + self.scan_cost += *scan_cost; + self.num_fetch_by_batches += *num_fetch_by_batches; + self.num_fetch_by_rows += *num_fetch_by_rows; + self.fetch_cost += *fetch_cost; + } + + /// Reports the metrics if scan_cost exceeds 10ms and resets them. + pub(crate) fn maybe_report(&mut self, reporter: &Option>) { + if self.scan_cost.as_millis() > 10 + && let Some(r) = reporter + { + r.report(self); + } + } +} + /// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`] // Ports from https://github.com/apache/datafusion/blob/49.0.0/datafusion/physical-plan/src/sorts/builder.rs // Adds the `take_remaining_rows()` method. diff --git a/src/mito2/src/read/merge.rs b/src/mito2/src/read/merge.rs deleted file mode 100644 index 0470e4b01a..0000000000 --- a/src/mito2/src/read/merge.rs +++ /dev/null @@ -1,982 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Merge reader implementation. - -use std::cmp::Ordering; -use std::collections::BinaryHeap; -use std::sync::Arc; -use std::time::{Duration, Instant}; -use std::{fmt, mem}; - -use async_trait::async_trait; -use common_telemetry::debug; - -use crate::error::Result; -use crate::memtable::BoxedBatchIterator; -use crate::metrics::READ_STAGE_ELAPSED; -use crate::read::{Batch, BatchReader, BoxedBatchReader, Source}; - -/// Trait for reporting merge metrics. -pub trait MergeMetricsReport: Send + Sync { - /// Reports and resets the metrics. - fn report(&self, metrics: &mut MergeMetrics); -} - -/// Reader to merge sorted batches. -/// -/// The merge reader merges [Batch]es from multiple sources that yield sorted batches. -/// 1. Batch is ordered by primary key, time index, sequence desc, op type desc (we can -/// ignore op type as sequence is already unique). -/// 2. Batches from sources **must** not be empty. -/// -/// The reader won't concatenate batches. Each batch returned by the reader also doesn't -/// contain duplicate rows. But the last (primary key, timestamp) of a batch may be the same -/// as the first one in the next batch. -pub struct MergeReader { - /// Holds [Node]s whose key range of current batch **is** overlapped with the merge window. - /// Each node yields batches from a `source`. - /// - /// [Node] in this heap **must** not be empty. A `merge window` is the (primary key, timestamp) - /// range of the **root node** in the `hot` heap. - hot: BinaryHeap, - /// Holds `Node` whose key range of current batch **isn't** overlapped with the merge window. - /// - /// `Node` in this heap **must** not be empty. - cold: BinaryHeap, - /// Batch to output. - output_batch: Option, - /// Local metrics. - metrics: MergeMetrics, - /// Optional metrics reporter. - metrics_reporter: Option>, -} - -#[async_trait] -impl BatchReader for MergeReader { - async fn next_batch(&mut self) -> Result> { - let start = Instant::now(); - while !self.hot.is_empty() && self.output_batch.is_none() { - if self.hot.len() == 1 { - // No need to do merge sort if only one batch in the hot heap. - self.fetch_batch_from_hottest().await?; - self.metrics.num_fetch_by_batches += 1; - } else { - // We could only fetch rows that less than the next node from the hottest node. - self.fetch_rows_from_hottest().await?; - self.metrics.num_fetch_by_rows += 1; - } - } - - if let Some(batch) = self.output_batch.take() { - self.metrics.scan_cost += start.elapsed(); - self.metrics.maybe_report(&self.metrics_reporter); - Ok(Some(batch)) - } else { - // Nothing fetched. - self.metrics.scan_cost += start.elapsed(); - self.metrics.maybe_report(&self.metrics_reporter); - Ok(None) - } - } -} - -impl Drop for MergeReader { - fn drop(&mut self) { - debug!("Merge reader finished, metrics: {:?}", self.metrics); - - READ_STAGE_ELAPSED - .with_label_values(&["merge"]) - .observe(self.metrics.scan_cost.as_secs_f64()); - READ_STAGE_ELAPSED - .with_label_values(&["merge_fetch"]) - .observe(self.metrics.fetch_cost.as_secs_f64()); - - // Report any remaining metrics. - if let Some(reporter) = &self.metrics_reporter { - reporter.report(&mut self.metrics); - } - } -} - -impl MergeReader { - /// Creates and initializes a new [MergeReader]. - pub async fn new( - sources: Vec, - metrics_reporter: Option>, - ) -> Result { - let start = Instant::now(); - let mut metrics = MergeMetrics::default(); - - let mut cold = BinaryHeap::with_capacity(sources.len()); - let hot = BinaryHeap::with_capacity(sources.len()); - for source in sources { - let node = Node::new(source, &mut metrics).await?; - if !node.is_eof() { - // Ensure `cold` don't have eof nodes. - cold.push(node); - } - } - - let mut reader = MergeReader { - hot, - cold, - output_batch: None, - metrics, - metrics_reporter, - }; - // Initializes the reader. - reader.refill_hot(); - - let elapsed = start.elapsed(); - reader.metrics.init_cost += elapsed; - reader.metrics.scan_cost += elapsed; - Ok(reader) - } - - /// Moves nodes in `cold` heap, whose key range is overlapped with current merge - /// window to `hot` heap. - fn refill_hot(&mut self) { - while !self.cold.is_empty() { - if let Some(merge_window) = self.hot.peek() { - let warmest = self.cold.peek().unwrap(); - if warmest.is_behind(merge_window) { - // if the warmest node in the `cold` heap is totally after the - // `merge_window`, then no need to add more nodes into the `hot` - // heap for merge sorting. - break; - } - } - - let warmest = self.cold.pop().unwrap(); - self.hot.push(warmest); - } - } - - /// Fetches one batch from the hottest node. - async fn fetch_batch_from_hottest(&mut self) -> Result<()> { - assert_eq!(1, self.hot.len()); - - let mut hottest = self.hot.pop().unwrap(); - let batch = hottest.fetch_batch(&mut self.metrics).await?; - Self::maybe_output_batch(batch, &mut self.output_batch)?; - self.reheap(hottest) - } - - /// Fetches non-duplicated rows from the hottest node. - async fn fetch_rows_from_hottest(&mut self) -> Result<()> { - // Safety: `fetch_batches_to_output()` ensures the hot heap has more than 1 element. - // Pop hottest node. - let mut top_node = self.hot.pop().unwrap(); - let top = top_node.current_batch(); - // Min timestamp and its sequence in the next batch. - let next_min_ts = { - let next_node = self.hot.peek().unwrap(); - let next = next_node.current_batch(); - // top and next have overlapping rows so they must have same primary keys. - debug_assert_eq!(top.primary_key(), next.primary_key()); - // Safety: Batches in the heap is not empty, so we can use unwrap here. - next.first_timestamp().unwrap() - }; - - // Safety: Batches in the heap is not empty, so we can use unwrap here. - let timestamps = top.timestamps_native().unwrap(); - // Binary searches the timestamp in the top batch. - // Safety: Batches should have the same timestamp resolution so we can compare the native - // value directly. - let duplicate_pos = match timestamps.binary_search(&next_min_ts.value()) { - Ok(pos) => pos, - Err(pos) => { - // No duplicate timestamp. Outputs timestamp before `pos`. - Self::maybe_output_batch(top.slice(0, pos), &mut self.output_batch)?; - top_node.skip_rows(pos, &mut self.metrics).await?; - return self.reheap(top_node); - } - }; - - // No need to remove duplicate timestamps. - let output_end = if duplicate_pos == 0 { - // If the first timestamp of the top node is duplicate. We can simply return the first row - // as the heap ensure it is the one with largest sequence. - 1 - } else { - // We don't know which one has the larger sequence so we use the range before - // the duplicate pos. - duplicate_pos - }; - Self::maybe_output_batch(top.slice(0, output_end), &mut self.output_batch)?; - top_node.skip_rows(output_end, &mut self.metrics).await?; - self.reheap(top_node) - } - - /// Push the node popped from `hot` back to a proper heap. - fn reheap(&mut self, node: Node) -> Result<()> { - if node.is_eof() { - // If the node is EOF, don't put it into the heap again. - // The merge window would be updated, need to refill the hot heap. - self.refill_hot(); - } else { - // Find a proper heap for this node. - let node_is_cold = if let Some(hottest) = self.hot.peek() { - // If key range of this node is behind the hottest node's then we can - // push it to the cold heap. Otherwise we should push it to the hot heap. - node.is_behind(hottest) - } else { - // The hot heap is empty, but we don't known whether the current - // batch of this node is still the hottest. - true - }; - - if node_is_cold { - self.cold.push(node); - } else { - self.hot.push(node); - } - // Anyway, the merge window has been changed, we need to refill the hot heap. - self.refill_hot(); - } - - Ok(()) - } - - /// If `filter_deleted` is set to true, removes deleted entries and sets the `batch` to the `output_batch`. - /// - /// Ignores the `batch` if it is empty. - fn maybe_output_batch(batch: Batch, output_batch: &mut Option) -> Result<()> { - debug_assert!(output_batch.is_none()); - if batch.is_empty() { - return Ok(()); - } - *output_batch = Some(batch); - - Ok(()) - } -} - -/// Builder to build and initialize a [MergeReader]. -#[derive(Default)] -pub struct MergeReaderBuilder { - /// Input sources. - /// - /// All source must yield batches with the same schema. - sources: Vec, - /// Optional metrics reporter. - metrics_reporter: Option>, -} - -impl MergeReaderBuilder { - /// Returns an empty builder. - pub fn new() -> MergeReaderBuilder { - MergeReaderBuilder::default() - } - - /// Creates a builder from sources. - pub fn from_sources(sources: Vec) -> MergeReaderBuilder { - MergeReaderBuilder { - sources, - metrics_reporter: None, - } - } - - /// Pushes a batch reader to sources. - pub fn push_batch_reader(&mut self, reader: BoxedBatchReader) -> &mut Self { - self.sources.push(Source::Reader(reader)); - self - } - - /// Pushes a batch iterator to sources. - pub fn push_batch_iter(&mut self, iter: BoxedBatchIterator) -> &mut Self { - self.sources.push(Source::Iter(iter)); - self - } - - /// Sets the metrics reporter. - pub fn with_metrics_reporter( - &mut self, - reporter: Option>, - ) -> &mut Self { - self.metrics_reporter = reporter; - self - } - - /// Builds and initializes the reader, then resets the builder. - pub async fn build(&mut self) -> Result { - let sources = mem::take(&mut self.sources); - let metrics_reporter = self.metrics_reporter.take(); - MergeReader::new(sources, metrics_reporter).await - } -} - -/// Metrics for the merge reader. -#[derive(Default)] -pub struct MergeMetrics { - /// Cost to initialize the reader. - pub(crate) init_cost: Duration, - /// Total scan cost of the reader. - pub(crate) scan_cost: Duration, - /// Number of times to fetch batches. - pub(crate) num_fetch_by_batches: usize, - /// Number of times to fetch rows. - pub(crate) num_fetch_by_rows: usize, - /// Cost to fetch batches from sources. - pub(crate) fetch_cost: Duration, -} - -impl fmt::Debug for MergeMetrics { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Skip output if scan_cost is zero - if self.scan_cost.is_zero() { - return write!(f, "{{}}"); - } - - write!(f, r#"{{"scan_cost":"{:?}""#, self.scan_cost)?; - - if !self.init_cost.is_zero() { - write!(f, r#", "init_cost":"{:?}""#, self.init_cost)?; - } - if self.num_fetch_by_batches > 0 { - write!( - f, - r#", "num_fetch_by_batches":{}"#, - self.num_fetch_by_batches - )?; - } - if self.num_fetch_by_rows > 0 { - write!(f, r#", "num_fetch_by_rows":{}"#, self.num_fetch_by_rows)?; - } - if !self.fetch_cost.is_zero() { - write!(f, r#", "fetch_cost":"{:?}""#, self.fetch_cost)?; - } - - write!(f, "}}") - } -} - -impl MergeMetrics { - /// Merges metrics from another MergeMetrics instance. - pub(crate) fn merge(&mut self, other: &MergeMetrics) { - let MergeMetrics { - init_cost, - scan_cost, - num_fetch_by_batches, - num_fetch_by_rows, - fetch_cost, - } = other; - - self.init_cost += *init_cost; - self.scan_cost += *scan_cost; - self.num_fetch_by_batches += *num_fetch_by_batches; - self.num_fetch_by_rows += *num_fetch_by_rows; - self.fetch_cost += *fetch_cost; - } - - /// Reports the metrics if scan_cost exceeds 10ms and resets them. - pub(crate) fn maybe_report(&mut self, reporter: &Option>) { - if self.scan_cost.as_millis() > 10 - && let Some(r) = reporter - { - r.report(self); - } - } -} - -/// A `Node` represent an individual input data source to be merged. -struct Node { - /// Data source of this `Node`. - source: Source, - /// Current batch to be read. The node ensures the batch is not empty. - /// - /// `None` means the `source` has reached EOF. - current_batch: Option, -} - -impl Node { - /// Initialize a node. - /// - /// It tries to fetch one batch from the `source`. - async fn new(mut source: Source, metrics: &mut MergeMetrics) -> Result { - // Ensures batch is not empty. - let start = Instant::now(); - let current_batch = source.next_batch().await?.map(CompareFirst); - metrics.fetch_cost += start.elapsed(); - - Ok(Node { - source, - current_batch, - }) - } - - /// Returns whether the node still has batch to read. - fn is_eof(&self) -> bool { - self.current_batch.is_none() - } - - /// Returns the primary key of current batch. - /// - /// # Panics - /// Panics if the node has reached EOF. - fn primary_key(&self) -> &[u8] { - self.current_batch().primary_key() - } - - /// Returns current batch. - /// - /// # Panics - /// Panics if the node has reached EOF. - fn current_batch(&self) -> &Batch { - &self.current_batch.as_ref().unwrap().0 - } - - /// Returns current batch and fetches next batch - /// from the source. - /// - /// # Panics - /// Panics if the node has reached EOF. - async fn fetch_batch(&mut self, metrics: &mut MergeMetrics) -> Result { - let current = self.current_batch.take().unwrap(); - let start = Instant::now(); - // Ensures batch is not empty. - self.current_batch = self.source.next_batch().await?.map(CompareFirst); - metrics.fetch_cost += start.elapsed(); - Ok(current.0) - } - - /// Returns true if the key range of current batch in `self` is behind (exclusive) current - /// batch in `other`. - /// - /// # Panics - /// Panics if either `self` or `other` is EOF. - fn is_behind(&self, other: &Node) -> bool { - debug_assert!(!self.current_batch().is_empty()); - debug_assert!(!other.current_batch().is_empty()); - - // We only compare pk and timestamp so nodes in the cold - // heap don't have overlapping timestamps with the hottest node - // in the hot heap. - self.primary_key().cmp(other.primary_key()).then_with(|| { - self.current_batch() - .first_timestamp() - .cmp(&other.current_batch().last_timestamp()) - }) == Ordering::Greater - } - - /// Skips first `num_to_skip` rows from node's current batch. If current batch is empty it fetches - /// next batch from the node. - /// - /// # Panics - /// Panics if the node is EOF. - async fn skip_rows(&mut self, num_to_skip: usize, metrics: &mut MergeMetrics) -> Result<()> { - let batch = self.current_batch(); - debug_assert!(batch.num_rows() >= num_to_skip); - - let remaining = batch.num_rows() - num_to_skip; - if remaining == 0 { - // Nothing remains, we need to fetch next batch to ensure the batch is not empty. - self.fetch_batch(metrics).await?; - } else { - debug_assert!(!batch.is_empty()); - self.current_batch = Some(CompareFirst(batch.slice(num_to_skip, remaining))); - } - - Ok(()) - } -} - -impl PartialEq for Node { - fn eq(&self, other: &Node) -> bool { - self.current_batch == other.current_batch - } -} - -impl Eq for Node {} - -impl PartialOrd for Node { - fn partial_cmp(&self, other: &Node) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for Node { - fn cmp(&self, other: &Node) -> Ordering { - // The std binary heap is a max heap, but we want the nodes are ordered in - // ascend order, so we compare the nodes in reverse order. - other.current_batch.cmp(&self.current_batch) - } -} - -/// Type to compare [Batch] by first row. -/// -/// It ignores op type as sequence is enough to distinguish different rows. -struct CompareFirst(Batch); - -impl PartialEq for CompareFirst { - fn eq(&self, other: &Self) -> bool { - self.0.primary_key() == other.0.primary_key() - && self.0.first_timestamp() == other.0.first_timestamp() - && self.0.first_sequence() == other.0.first_sequence() - } -} - -impl Eq for CompareFirst {} - -impl PartialOrd for CompareFirst { - fn partial_cmp(&self, other: &CompareFirst) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for CompareFirst { - /// Compares by primary key, time index, sequence desc. - fn cmp(&self, other: &CompareFirst) -> Ordering { - self.0 - .primary_key() - .cmp(other.0.primary_key()) - .then_with(|| self.0.first_timestamp().cmp(&other.0.first_timestamp())) - .then_with(|| other.0.first_sequence().cmp(&self.0.first_sequence())) - } -} - -#[cfg(test)] -mod tests { - use api::v1::OpType; - - use super::*; - use crate::test_util::{VecBatchReader, check_reader_result, new_batch}; - - #[tokio::test] - async fn test_merge_reader_empty() { - let mut reader = MergeReaderBuilder::new().build().await.unwrap(); - assert!(reader.next_batch().await.unwrap().is_none()); - assert!(reader.next_batch().await.unwrap().is_none()); - } - - #[tokio::test] - async fn test_merge_non_overlapping() { - let reader1 = VecBatchReader::new(&[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch( - b"k1", - &[7, 8], - &[17, 18], - &[OpType::Put, OpType::Delete], - &[27, 28], - ), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - ]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[4, 5], - &[14, 15], - &[OpType::Put, OpType::Put], - &[24, 25], - )]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch( - b"k1", - &[4, 5], - &[14, 15], - &[OpType::Put, OpType::Put], - &[24, 25], - ), - new_batch( - b"k1", - &[7, 8], - &[17, 18], - &[OpType::Put, OpType::Delete], - &[27, 28], - ), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_reheap_hot() { - let reader1 = VecBatchReader::new(&[ - new_batch( - b"k1", - &[1, 3], - &[10, 10], - &[OpType::Put, OpType::Put], - &[21, 23], - ), - new_batch(b"k2", &[3], &[10], &[OpType::Put], &[23]), - ]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[2, 4], - &[11, 11], - &[OpType::Put, OpType::Put], - &[32, 34], - )]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[1], &[10], &[OpType::Put], &[21]), - new_batch(b"k1", &[2], &[11], &[OpType::Put], &[32]), - new_batch(b"k1", &[3], &[10], &[OpType::Put], &[23]), - new_batch(b"k1", &[4], &[11], &[OpType::Put], &[34]), - new_batch(b"k2", &[3], &[10], &[OpType::Put], &[23]), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_overlapping() { - let reader1 = VecBatchReader::new(&[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch( - b"k1", - &[4, 5], - &[14, 15], - // This override 4 and deletes 5. - &[OpType::Put, OpType::Delete], - &[24, 25], - ), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - // This delete 2. - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - ]); - let reader2 = VecBatchReader::new(&[ - new_batch( - b"k1", - &[3, 4, 5], - &[10, 10, 10], - &[OpType::Put, OpType::Put, OpType::Put], - &[33, 34, 35], - ), - new_batch( - b"k2", - &[1, 10], - &[11, 20], - &[OpType::Put, OpType::Put], - &[21, 30], - ), - ]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch(b"k1", &[3], &[10], &[OpType::Put], &[33]), - new_batch(b"k1", &[4], &[14], &[OpType::Put], &[24]), - new_batch(b"k1", &[4], &[10], &[OpType::Put], &[34]), - new_batch(b"k1", &[5], &[15], &[OpType::Delete], &[25]), - new_batch(b"k1", &[5], &[10], &[OpType::Put], &[35]), - new_batch(b"k2", &[1], &[11], &[OpType::Put], &[21]), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - new_batch(b"k2", &[10], &[20], &[OpType::Put], &[30]), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_deleted() { - let reader1 = VecBatchReader::new(&[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Delete, OpType::Delete], - &[21, 22], - ), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - ]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[4, 5], - &[14, 15], - &[OpType::Delete, OpType::Delete], - &[24, 25], - )]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Delete, OpType::Delete], - &[21, 22], - ), - new_batch( - b"k1", - &[4, 5], - &[14, 15], - &[OpType::Delete, OpType::Delete], - &[24, 25], - ), - new_batch( - b"k2", - &[2, 3], - &[12, 13], - &[OpType::Delete, OpType::Put], - &[22, 23], - ), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_next_node_empty() { - let reader1 = VecBatchReader::new(&[new_batch( - b"k1", - &[1, 2], - &[11, 12], - &[OpType::Put, OpType::Put], - &[21, 22], - )]); - let reader2 = VecBatchReader::new(&[new_batch(b"k1", &[1], &[10], &[OpType::Put], &[33])]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[1], &[11], &[OpType::Put], &[21]), - new_batch(b"k1", &[1], &[10], &[OpType::Put], &[33]), - new_batch(b"k1", &[2], &[12], &[OpType::Put], &[22]), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_top_node_empty() { - let reader1 = VecBatchReader::new(&[new_batch( - b"k1", - &[1, 2], - &[10, 10], - &[OpType::Put, OpType::Put], - &[21, 22], - )]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[2, 3], - &[11, 11], - &[OpType::Put, OpType::Put], - &[32, 33], - )]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[1], &[10], &[OpType::Put], &[21]), - new_batch(b"k1", &[2], &[11], &[OpType::Put], &[32]), - new_batch(b"k1", &[2], &[10], &[OpType::Put], &[22]), - new_batch(b"k1", &[3], &[11], &[OpType::Put], &[33]), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_large_range() { - let reader1 = VecBatchReader::new(&[new_batch( - b"k1", - &[1, 10], - &[10, 10], - &[OpType::Put, OpType::Put], - &[21, 30], - )]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[1, 20], - &[11, 11], - &[OpType::Put, OpType::Put], - &[31, 40], - )]); - // The hot heap have a node that doesn't have duplicate - // timestamps. - let reader3 = VecBatchReader::new(&[new_batch( - b"k1", - &[6, 8], - &[11, 11], - &[OpType::Put, OpType::Put], - &[36, 38], - )]); - let mut reader = MergeReaderBuilder::new() - .push_batch_reader(Box::new(reader1)) - .push_batch_iter(Box::new(reader2)) - .push_batch_reader(Box::new(reader3)) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[1], &[11], &[OpType::Put], &[31]), - new_batch(b"k1", &[1], &[10], &[OpType::Put], &[21]), - new_batch( - b"k1", - &[6, 8], - &[11, 11], - &[OpType::Put, OpType::Put], - &[36, 38], - ), - new_batch(b"k1", &[10], &[10], &[OpType::Put], &[30]), - new_batch(b"k1", &[20], &[11], &[OpType::Put], &[40]), - ], - ) - .await; - } - - #[tokio::test] - async fn test_merge_many_duplicates() { - let mut builder = MergeReaderBuilder::new(); - for i in 0..10 { - let batches: Vec<_> = (0..8) - .map(|ts| new_batch(b"k1", &[ts], &[i], &[OpType::Put], &[100])) - .collect(); - let reader = VecBatchReader::new(&batches); - builder.push_batch_reader(Box::new(reader)); - } - let mut reader = builder.build().await.unwrap(); - let mut expect = Vec::with_capacity(80); - for ts in 0..8 { - for i in 0..10 { - let batch = new_batch(b"k1", &[ts], &[9 - i], &[OpType::Put], &[100]); - expect.push(batch); - } - } - check_reader_result(&mut reader, &expect).await; - } - - #[tokio::test] - async fn test_merge_keep_duplicate() { - let reader1 = VecBatchReader::new(&[new_batch( - b"k1", - &[1, 2], - &[10, 10], - &[OpType::Put, OpType::Put], - &[21, 22], - )]); - let reader2 = VecBatchReader::new(&[new_batch( - b"k1", - &[2, 3], - &[11, 11], - &[OpType::Put, OpType::Put], - &[32, 33], - )]); - let sources = vec![ - Source::Reader(Box::new(reader1)), - Source::Iter(Box::new(reader2)), - ]; - let mut reader = MergeReaderBuilder::from_sources(sources) - .build() - .await - .unwrap(); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[1], &[10], &[OpType::Put], &[21]), - new_batch(b"k1", &[2], &[11], &[OpType::Put], &[32]), - new_batch(b"k1", &[2], &[10], &[OpType::Put], &[22]), - new_batch(b"k1", &[3], &[11], &[OpType::Put], &[33]), - ], - ) - .await; - } -} diff --git a/src/mito2/src/read/plain_batch.rs b/src/mito2/src/read/plain_batch.rs deleted file mode 100644 index f22b6688d6..0000000000 --- a/src/mito2/src/read/plain_batch.rs +++ /dev/null @@ -1,505 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Batch without an encoded primary key. - -use std::collections::HashMap; -use std::sync::Arc; - -use api::v1::OpType; -use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt8Array, UInt64Array}; -use datatypes::arrow::compute::filter_record_batch; -use datatypes::arrow::datatypes::SchemaRef; -use datatypes::arrow::record_batch::RecordBatch; -use snafu::{OptionExt, ResultExt}; -use store_api::metadata::{ColumnMetadata, RegionMetadata}; -use store_api::storage::{RegionId, SequenceNumber}; - -use crate::error::{ - ComputeArrowSnafu, CreateDefaultSnafu, InvalidRequestSnafu, NewRecordBatchSnafu, Result, - UnexpectedSnafu, -}; - -/// Number of columns that have fixed positions. -/// -/// Contains all internal columns. -pub(crate) const PLAIN_FIXED_POS_COLUMN_NUM: usize = 2; - -/// [PlainBatch] represents a batch of rows. -/// It is a wrapper around [RecordBatch]. -/// -/// The columns order is the same as the order of the columns read from the SST. -/// It always contains two internal columns now. We may change modify this behavior -/// in the future. -#[derive(Debug)] -pub struct PlainBatch { - /// The original record batch. - record_batch: RecordBatch, -} - -impl PlainBatch { - /// Creates a new [PlainBatch] from a [RecordBatch]. - pub fn new(record_batch: RecordBatch) -> Self { - assert!( - record_batch.num_columns() >= 2, - "record batch missing internal columns, num_columns: {}", - record_batch.num_columns() - ); - - Self { record_batch } - } - - /// Returns a new [PlainBatch] with the given columns. - pub fn with_new_columns(&self, columns: Vec) -> Result { - let record_batch = RecordBatch::try_new(self.record_batch.schema(), columns) - .context(NewRecordBatchSnafu)?; - Ok(Self::new(record_batch)) - } - - /// Returns the number of columns in the batch. - pub fn num_columns(&self) -> usize { - self.record_batch.num_columns() - } - - /// Returns the number of rows in the batch. - pub fn num_rows(&self) -> usize { - self.record_batch.num_rows() - } - - /// Returns true if the batch is empty. - pub fn is_empty(&self) -> bool { - self.num_rows() == 0 - } - - /// Returns all columns. - pub fn columns(&self) -> &[ArrayRef] { - self.record_batch.columns() - } - - /// Returns the array of column at index `idx`. - pub fn column(&self, idx: usize) -> &ArrayRef { - self.record_batch.column(idx) - } - - /// Returns the slice of internal columns. - pub fn internal_columns(&self) -> &[ArrayRef] { - &self.record_batch.columns()[self.record_batch.num_columns() - PLAIN_FIXED_POS_COLUMN_NUM..] - } - - /// Returns the inner record batch. - pub fn as_record_batch(&self) -> &RecordBatch { - &self.record_batch - } - - /// Converts this batch into a record batch. - pub fn into_record_batch(self) -> RecordBatch { - self.record_batch - } - - /// Filters this batch by the boolean array. - pub fn filter(&self, predicate: &BooleanArray) -> Result { - let record_batch = - filter_record_batch(&self.record_batch, predicate).context(ComputeArrowSnafu)?; - Ok(Self::new(record_batch)) - } - - /// Returns the column index of the sequence column. - #[allow(dead_code)] - pub(crate) fn sequence_column_index(&self) -> usize { - self.record_batch.num_columns() - PLAIN_FIXED_POS_COLUMN_NUM - } -} - -/// Helper struct to fill default values and internal columns. -pub struct ColumnFiller<'a> { - /// Region metadata information - metadata: &'a RegionMetadata, - /// Schema for the output record batch - schema: SchemaRef, - /// Map of column names to indices in the input record batch - name_to_index: HashMap, -} - -impl<'a> ColumnFiller<'a> { - /// Creates a new ColumnFiller - /// The `schema` is the sst schema of the `metadata`. - pub fn new( - metadata: &'a RegionMetadata, - schema: SchemaRef, - record_batch: &RecordBatch, - ) -> Self { - debug_assert_eq!(metadata.column_metadatas.len() + 2, schema.fields().len()); - - // Pre-construct the name to index map - let name_to_index: HashMap<_, _> = record_batch - .schema() - .fields() - .iter() - .enumerate() - .map(|(i, field)| (field.name().clone(), i)) - .collect(); - - Self { - metadata, - schema, - name_to_index, - } - } - - /// Fills default values and internal columns for a [RecordBatch]. - pub fn fill_missing_columns( - &self, - record_batch: &RecordBatch, - sequence: SequenceNumber, - op_type: OpType, - ) -> Result { - let num_rows = record_batch.num_rows(); - let mut new_columns = - Vec::with_capacity(record_batch.num_columns() + PLAIN_FIXED_POS_COLUMN_NUM); - - // Fills default values. - // Implementation based on `WriteRequest::fill_missing_columns()`. - for column in &self.metadata.column_metadatas { - let array = match self.name_to_index.get(&column.column_schema.name) { - Some(index) => record_batch.column(*index).clone(), - None => match op_type { - OpType::Put => { - // For put requests, we use the default value from column schema. - fill_column_put_default(self.metadata.region_id, column, num_rows)? - } - OpType::Delete => { - // For delete requests, we need default value for padding. - fill_column_delete_default(column, num_rows)? - } - }, - }; - - new_columns.push(array); - } - - // Adds internal columns. - // Adds the sequence number. - let sequence_array = Arc::new(UInt64Array::from(vec![sequence; num_rows])); - // Adds the op type. - let op_type_array = Arc::new(UInt8Array::from(vec![op_type as u8; num_rows])); - new_columns.push(sequence_array); - new_columns.push(op_type_array); - - RecordBatch::try_new(self.schema.clone(), new_columns).context(NewRecordBatchSnafu) - } -} - -fn fill_column_put_default( - region_id: RegionId, - column: &ColumnMetadata, - num_rows: usize, -) -> Result { - if column.column_schema.is_default_impure() { - return UnexpectedSnafu { - reason: format!( - "unexpected impure default value with region_id: {}, column: {}, default_value: {:?}", - region_id, - column.column_schema.name, - column.column_schema.default_constraint(), - ), - } - .fail(); - } - let vector = column - .column_schema - .create_default_vector(num_rows) - .context(CreateDefaultSnafu { - region_id, - column: &column.column_schema.name, - })? - // This column doesn't have default value. - .with_context(|| InvalidRequestSnafu { - region_id, - reason: format!( - "column {} does not have default value", - column.column_schema.name - ), - })?; - Ok(vector.to_arrow_array()) -} - -fn fill_column_delete_default(column: &ColumnMetadata, num_rows: usize) -> Result { - // For delete requests, we need a default value for padding - let vector = column - .column_schema - .create_default_vector_for_padding(num_rows); - Ok(vector.to_arrow_array()) -} - -#[cfg(test)] -mod tests { - use api::v1::SemanticType; - use datatypes::arrow::array::{ - Float64Array, Int32Array, StringArray, TimestampMillisecondArray, - }; - use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; - use datatypes::schema::ColumnSchema; - use datatypes::schema::constraint::ColumnDefaultConstraint; - use datatypes::value::Value; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; - use store_api::storage::{ConcreteDataType, RegionId}; - - use super::*; - use crate::sst::to_plain_sst_arrow_schema; - - /// Creates a test region metadata with schema: k0(string), ts(timestamp), v1(float64) - fn create_test_region_metadata() -> RegionMetadata { - let mut builder = RegionMetadataBuilder::new(RegionId::new(100, 200)); - builder - // Add string key column - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), false) - .with_default_constraint(None) - .unwrap(), - semantic_type: SemanticType::Tag, - column_id: 0, - }) - // Add timestamp column - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true) - .with_default_constraint(None) - .unwrap(), - semantic_type: SemanticType::Timestamp, - column_id: 1, - }) - // Add float value column with default - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("v1", ConcreteDataType::float64_datatype(), true) - .with_default_constraint(Some(ColumnDefaultConstraint::Value(Value::Float64( - datatypes::value::OrderedFloat::from(42.0), - )))) - .unwrap(), - semantic_type: SemanticType::Field, - column_id: 2, - }) - .primary_key(vec![0]); - - builder.build().unwrap() - } - - #[test] - fn test_column_filler_put() { - let region_metadata = create_test_region_metadata(); - let output_schema = to_plain_sst_arrow_schema(®ion_metadata); - - // Create input record batch with only k0 and ts columns (v1 is missing) - let input_schema = Arc::new(Schema::new(vec![ - Field::new("k0", DataType::Utf8, false), - Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - ), - ])); - - let k0_values: ArrayRef = Arc::new(StringArray::from(vec!["key1", "key2"])); - let ts_values: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1000, 2000])); - - let input_batch = - RecordBatch::try_new(input_schema, vec![k0_values.clone(), ts_values.clone()]).unwrap(); - - // Create column filler - let filler = ColumnFiller::new(®ion_metadata, output_schema.clone(), &input_batch); - - // Fill missing columns with OpType::Put - let result = filler - .fill_missing_columns(&input_batch, 100, OpType::Put) - .unwrap(); - - // Verify the result - // Create an expected record batch to compare against - let expected_columns = vec![ - k0_values.clone(), - ts_values.clone(), - Arc::new(Float64Array::from(vec![42.0, 42.0])), - Arc::new(UInt64Array::from(vec![100, 100])), - Arc::new(UInt8Array::from(vec![OpType::Put as u8, OpType::Put as u8])), - ]; - let expected_batch = RecordBatch::try_new(output_schema.clone(), expected_columns).unwrap(); - assert_eq!(expected_batch, result); - } - - #[test] - fn test_column_filler_delete() { - let region_metadata = create_test_region_metadata(); - let output_schema = to_plain_sst_arrow_schema(®ion_metadata); - - // Create input record batch with only k0 and ts columns (v1 is missing) - let input_schema = Arc::new(Schema::new(vec![ - Field::new("k0", DataType::Utf8, false), - Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - ), - ])); - - let k0_values: ArrayRef = Arc::new(StringArray::from(vec!["key1", "key2"])); - let ts_values: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1000, 2000])); - - let input_batch = - RecordBatch::try_new(input_schema, vec![k0_values.clone(), ts_values.clone()]).unwrap(); - - // Create column filler - let filler = ColumnFiller::new(®ion_metadata, output_schema.clone(), &input_batch); - - // Fill missing columns with OpType::Delete - let result = filler - .fill_missing_columns(&input_batch, 200, OpType::Delete) - .unwrap(); - - // Verify the result by creating an expected record batch to compare against - let v1_default = Arc::new(Float64Array::from(vec![None, None])); - let expected_columns = vec![ - k0_values.clone(), - ts_values.clone(), - v1_default, - Arc::new(UInt64Array::from(vec![200, 200])), - Arc::new(UInt8Array::from(vec![ - OpType::Delete as u8, - OpType::Delete as u8, - ])), - ]; - let expected_batch = RecordBatch::try_new(output_schema.clone(), expected_columns).unwrap(); - assert_eq!(expected_batch, result); - } - - fn create_test_record_batch() -> RecordBatch { - let schema = Arc::new(Schema::new(vec![ - Field::new("col1", DataType::Int32, false), - Field::new("col2", DataType::Utf8, false), - Field::new(SEQUENCE_COLUMN_NAME, DataType::UInt64, false), - Field::new(OP_TYPE_COLUMN_NAME, DataType::UInt8, false), - ])); - - let col1 = Arc::new(Int32Array::from(vec![1, 2, 3])); - let col2 = Arc::new(StringArray::from(vec!["a", "b", "c"])); - let sequence = Arc::new(UInt64Array::from(vec![100, 101, 102])); - let op_type = Arc::new(UInt8Array::from(vec![1, 1, 1])); - - RecordBatch::try_new(schema, vec![col1, col2, sequence, op_type]).unwrap() - } - - #[test] - fn test_plain_batch_basic_methods() { - let record_batch = create_test_record_batch(); - let plain_batch = PlainBatch::new(record_batch.clone()); - - // Test basic properties - assert_eq!(plain_batch.num_columns(), 4); - assert_eq!(plain_batch.num_rows(), 3); - assert!(!plain_batch.is_empty()); - assert_eq!(plain_batch.columns().len(), 4); - - // Test internal columns access - let internal_columns = plain_batch.internal_columns(); - assert_eq!(internal_columns.len(), PLAIN_FIXED_POS_COLUMN_NUM); - assert_eq!(internal_columns[0].len(), 3); - assert_eq!(internal_columns[1].len(), 3); - - // Test column access - let col1 = plain_batch.column(0); - assert_eq!(col1.len(), 3); - assert_eq!( - col1.as_any().downcast_ref::().unwrap().value(0), - 1 - ); - - // Test sequence column index - assert_eq!(plain_batch.sequence_column_index(), 2); - - // Test to record batch. - assert_eq!(record_batch, *plain_batch.as_record_batch()); - assert_eq!(record_batch, plain_batch.into_record_batch()); - } - - #[test] - fn test_with_new_columns() { - let record_batch = create_test_record_batch(); - let plain_batch = PlainBatch::new(record_batch); - - // Create new columns - let col1 = Arc::new(Int32Array::from(vec![10, 20, 30])); - let col2 = Arc::new(StringArray::from(vec!["x", "y", "z"])); - let sequence = Arc::new(UInt64Array::from(vec![200, 201, 202])); - let op_type = Arc::new(UInt8Array::from(vec![0, 0, 0])); - - let new_batch = plain_batch - .with_new_columns(vec![col1, col2, sequence, op_type]) - .unwrap(); - - assert_eq!(new_batch.num_columns(), 4); - assert_eq!(new_batch.num_rows(), 3); - assert_eq!( - new_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 10 - ); - assert_eq!( - new_batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - "x" - ); - } - - #[test] - fn test_filter() { - let record_batch = create_test_record_batch(); - let plain_batch = PlainBatch::new(record_batch); - - // Create a predicate that selects the first and third rows - let predicate = BooleanArray::from(vec![true, false, true]); - - let filtered_batch = plain_batch.filter(&predicate).unwrap(); - - assert_eq!(filtered_batch.num_rows(), 2); - assert_eq!( - filtered_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 1 - ); - assert_eq!( - filtered_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(1), - 3 - ); - } -} diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index c447685822..f645e3dc26 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -60,7 +60,7 @@ use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; use crate::read::unordered_scan::UnorderedScan; -use crate::read::{Batch, BoxedRecordBatchStream, RecordBatch, Source}; +use crate::read::{BoxedRecordBatchStream, RecordBatch}; use crate::region::options::MergeMode; use crate::region::version::VersionRef; use crate::sst::file::FileHandle; @@ -1031,39 +1031,6 @@ impl ScanInput { self } - /// Scans sources in parallel. - /// - /// # Panics if the input doesn't allow parallel scan. - #[tracing::instrument( - skip(self, sources, semaphore), - fields( - region_id = %self.region_metadata().region_id, - source_count = sources.len() - ) - )] - pub(crate) fn create_parallel_sources( - &self, - sources: Vec, - semaphore: Arc, - channel_size: usize, - ) -> Result> { - if sources.len() <= 1 { - return Ok(sources); - } - - // Spawn a task for each source. - let sources = sources - .into_iter() - .map(|source| { - let (sender, receiver) = mpsc::channel(channel_size); - self.spawn_scan_task(source, semaphore.clone(), sender); - let stream = Box::pin(ReceiverStream::new(receiver)); - Source::Stream(stream) - }) - .collect(); - Ok(sources) - } - /// Builds memtable ranges to scan by `index`. pub(crate) fn build_mem_ranges(&self, index: RowGroupIndex) -> SmallVec<[MemtableRange; 2]> { let memtable = &self.memtables[index.index]; @@ -1173,49 +1140,6 @@ impl ScanInput { Ok(FileRangeBuilder::new(Arc::new(file_range_ctx), selection)) } - /// Scans the input source in another task and sends batches to the sender. - #[tracing::instrument( - skip(self, input, semaphore, sender), - fields(region_id = %self.region_metadata().region_id) - )] - pub(crate) fn spawn_scan_task( - &self, - mut input: Source, - semaphore: Arc, - sender: mpsc::Sender>, - ) { - let region_id = self.region_metadata().region_id; - let span = tracing::info_span!( - "ScanInput::parallel_scan_task", - region_id = %region_id, - stream_kind = "batch" - ); - common_runtime::spawn_global( - async move { - loop { - // We release the permit before sending result to avoid the task waiting on - // the channel with the permit held. - let maybe_batch = { - // Safety: We never close the semaphore. - let _permit = semaphore.acquire().await.unwrap(); - input.next_batch().await - }; - match maybe_batch { - Ok(Some(batch)) => { - let _ = sender.send(Ok(batch)).await; - } - Ok(None) => break, - Err(e) => { - let _ = sender.send(Err(e)).await; - break; - } - } - } - } - .instrument(span), - ); - } - /// Scans flat sources (RecordBatch streams) in parallel. /// /// # Panics if the input doesn't allow parallel scan. diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 597f592de6..8fc946b3d3 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -39,7 +39,7 @@ use crate::metrics::{ READ_ROWS_IN_ROW_GROUP_TOTAL, READ_ROWS_RETURN, READ_STAGE_ELAPSED, }; use crate::read::dedup::{DedupMetrics, DedupMetricsReport}; -use crate::read::merge::{MergeMetrics, MergeMetricsReport}; +use crate::read::flat_merge::{MergeMetrics, MergeMetricsReport}; use crate::read::pruner::PartitionPruner; use crate::read::range::{RangeMeta, RowGroupIndex}; use crate::read::scan_region::StreamContext; diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index 94bc1feea8..c769f78c6c 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -218,34 +218,6 @@ pub(crate) fn internal_fields() -> [FieldRef; 3] { ] } -/// Gets the arrow schema to store in parquet. -pub fn to_plain_sst_arrow_schema(metadata: &RegionMetadata) -> SchemaRef { - let fields = Fields::from_iter( - metadata - .schema - .arrow_schema() - .fields() - .iter() - .cloned() - .chain(plain_internal_fields()), - ); - - Arc::new(Schema::new(fields)) -} - -/// Fields for internal columns. -fn plain_internal_fields() -> [FieldRef; 2] { - // Internal columns are always not null. - [ - Arc::new(Field::new( - SEQUENCE_COLUMN_NAME, - ArrowDataType::UInt64, - false, - )), - Arc::new(Field::new(OP_TYPE_COLUMN_NAME, ArrowDataType::UInt8, false)), - ] -} - /// Gets the estimated number of series from record batches. /// /// This struct tracks the last timestamp value to detect series boundaries diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs index e9515030c0..84f15ad837 100644 --- a/src/mito2/src/test_util/sst_util.rs +++ b/src/mito2/src/test_util/sst_util.rs @@ -36,10 +36,10 @@ use store_api::metric_engine_consts::{ use store_api::storage::consts::ReservedColumnId; use store_api::storage::{FileId, RegionId}; -use crate::read::{Batch, FlatSource, Source}; +use crate::read::{Batch, FlatSource}; use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; -use crate::test_util::{VecBatchReader, new_batch_builder, new_noop_file_purger}; +use crate::test_util::{new_batch_builder, new_noop_file_purger}; /// Test region id. const REGION_ID: RegionId = RegionId::new(0, 0); @@ -190,12 +190,6 @@ pub fn new_sparse_primary_key( buffer } -/// Creates a [Source] from `batches`. -pub fn new_source(batches: &[Batch]) -> Source { - let reader = VecBatchReader::new(batches); - Source::Reader(Box::new(reader)) -} - /// Creates a SST file handle with provided file id pub fn sst_file_handle_with_file_id(file_id: FileId, start_ms: i64, end_ms: i64) -> FileHandle { let file_purger = new_noop_file_purger(); From 59021ce83b94e8457283f08ebb92898e00ae300f Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Fri, 10 Apr 2026 15:56:33 +0800 Subject: [PATCH 007/120] fix: using uint64 datatype for postgres prepared statement parameters (#7942) * feat: add support for decimal parameter type, remove string replacement fallback * chore: format * fix: add support for using unsigned bigint in postgres * chore: format toml * refactor: cleanup duplicated code * fix: rescale decimal --- Cargo.lock | 1 + src/servers/Cargo.toml | 1 + src/servers/src/postgres/handler.rs | 15 +-- src/servers/src/postgres/types.rs | 191 ++++++++++++++++++---------- tests-integration/tests/sql.rs | 52 ++++++++ 5 files changed, 184 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index edb8ce04d4..872095752b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12080,6 +12080,7 @@ dependencies = [ "regex", "reqwest", "rust-embed", + "rust_decimal", "rustls", "rustls-pemfile", "rustls-pki-types", diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 2d68f17699..46a51f1280 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -107,6 +107,7 @@ rand.workspace = true regex.workspace = true reqwest.workspace = true rust-embed = { version = "6.6", optional = true, features = ["debug-embed"] } +rust_decimal = { workspace = true, features = ["db-postgres"] } rustls = { workspace = true, default-features = false, features = ["ring", "logging", "std", "tls12"] } rustls-pemfile = "2.0" rustls-pki-types = "1.0" diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 2b84b3aa30..94363b06eb 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -456,16 +456,13 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner { .do_exec_plan(sql_plan.statement.clone(), plan, query_ctx.clone()) .await } else { - // manually replace variables in prepared statement when no - // logical_plan is generated. This happens when logical plan is not - // supported for certain statements. - let mut sql = sql_plan.query.clone(); - for i in 0..portal.parameter_len() { - sql = sql.replace(&format!("${}", i + 1), ¶meter_to_string(portal, i)?); - } - + // We won't replace params from statement manually any more. + // Newer version of datafusion can generate plan for SELECT/INSERT/UPDATE/DELETE. + // Only CREATE TABLE and others minor statements cannot generate sql plan, + // in this case, we assume these statements will not carry parameters + // and execute them directly. self.query_handler - .do_query(&sql, query_ctx.clone()) + .do_query(&sql_plan.query, query_ctx.clone()) .await .remove(0) }; diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index d4d15ef64a..203e477c6f 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -33,7 +33,7 @@ use datatypes::arrow::datatypes::DataType as ArrowDataType; use datatypes::json::JsonStructureSettings; use datatypes::prelude::{ConcreteDataType, Value}; use datatypes::schema::{Schema, SchemaRef}; -use datatypes::types::{IntervalType, TimestampType, jsonb_to_string}; +use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; use datatypes::value::StructValue; use futures::Stream; use pg_interval::Interval as PgInterval; @@ -43,6 +43,8 @@ use pgwire::api::results::FieldInfo; use pgwire::error::{PgWireError, PgWireResult}; use pgwire::types::format::FormatOptions as PgFormatOptions; use query::planner::DfLogicalPlanner; +use rust_decimal::Decimal; +use rust_decimal::prelude::ToPrimitive; use session::context::QueryContextRef; use snafu::ResultExt; @@ -293,11 +295,11 @@ pub(super) fn type_pg_to_gt(origin: &Type) -> Result { // Note that we only support a small amount of pg data types match origin { &Type::BOOL => Ok(ConcreteDataType::boolean_datatype()), - &Type::CHAR => Ok(ConcreteDataType::int8_datatype()), &Type::INT2 => Ok(ConcreteDataType::int16_datatype()), &Type::INT4 => Ok(ConcreteDataType::int32_datatype()), &Type::INT8 => Ok(ConcreteDataType::int64_datatype()), - &Type::VARCHAR | &Type::TEXT => Ok(ConcreteDataType::string_datatype()), + &Type::NUMERIC => Ok(ConcreteDataType::uint64_datatype()), + &Type::VARCHAR | &Type::CHAR | &Type::TEXT => Ok(ConcreteDataType::string_datatype()), &Type::TIMESTAMP | &Type::TIMESTAMPTZ => Ok(ConcreteDataType::timestamp_datatype( common_time::timestamp::TimeUnit::Millisecond, )), @@ -305,9 +307,6 @@ pub(super) fn type_pg_to_gt(origin: &Type) -> Result { &Type::TIME => Ok(ConcreteDataType::timestamp_datatype( common_time::timestamp::TimeUnit::Microsecond, )), - &Type::CHAR_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int8_datatype(), - ))), &Type::INT2_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( ConcreteDataType::int16_datatype(), ))), @@ -317,9 +316,12 @@ pub(super) fn type_pg_to_gt(origin: &Type) -> Result { &Type::INT8_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( ConcreteDataType::int64_datatype(), ))), - &Type::VARCHAR_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::string_datatype(), + &Type::NUMERIC_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( + ConcreteDataType::uint64_datatype(), ))), + &Type::VARCHAR_ARRAY | &Type::CHAR_ARRAY | &Type::TEXT_ARRAY => Ok( + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())), + ), _ => server_error::InternalSnafu { err_msg: format!("unimplemented datatype {origin:?}"), } @@ -327,63 +329,6 @@ pub(super) fn type_pg_to_gt(origin: &Type) -> Result { } } -pub(super) fn parameter_to_string(portal: &Portal, idx: usize) -> PgWireResult { - // the index is managed from portal's parameters count so it's safe to - // unwrap here. - let param_type = portal - .statement - .parameter_types - .get(idx) - .unwrap() - .as_ref() - .unwrap_or(&Type::UNKNOWN); - match param_type { - &Type::VARCHAR | &Type::TEXT => Ok(format!( - "'{}'", - portal - .parameter::(idx, param_type)? - .as_deref() - .unwrap_or("") - )), - &Type::BOOL => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::INT4 => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::INT8 => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::FLOAT4 => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::FLOAT8 => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::DATE => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.format("%Y-%m-%d").to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::TIMESTAMP => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.format("%Y-%m-%d %H:%M:%S%.6f").to_string()) - .unwrap_or_else(|| "".to_owned())), - &Type::INTERVAL => Ok(portal - .parameter::(idx, param_type)? - .map(|v| v.to_sql()) - .unwrap_or_else(|| "".to_owned())), - _ => Err(invalid_parameter_error( - "unsupported_parameter_type", - Some(param_type.to_string()), - )), - } -} - pub(super) fn invalid_parameter_error(msg: &str, detail: Option) -> PgWireError { let mut error_info = PgErrorCode::Ec22023.to_err_info(msg.to_string()); error_info.detail = detail; @@ -407,6 +352,17 @@ where } } +fn to_decimal_scalar_value(data: Option, ctype: &Decimal128Type) -> ScalarValue { + if let Some(data) = data { + let mut value = data; + value.rescale(ctype.scale() as u32); + + ScalarValue::Decimal128(Some(value.mantissa()), ctype.precision(), ctype.scale()) + } else { + ScalarValue::Decimal128(None, ctype.precision(), ctype.scale()) + } +} + pub(super) fn parameters_to_scalar_values( plan: &LogicalPlan, portal: &Portal, @@ -442,7 +398,7 @@ pub(super) fn parameters_to_scalar_values( }; let value = match &client_type { - &Type::VARCHAR | &Type::TEXT => { + &Type::VARCHAR | &Type::TEXT | &Type::CHAR => { let data = portal.parameter::(idx, &client_type)?; if let Some(server_type) = &server_type { match server_type { @@ -558,6 +514,24 @@ pub(super) fn parameters_to_scalar_values( ScalarValue::Int64(data) } } + &Type::NUMERIC => { + let data = portal.parameter::(idx, &client_type)?; + match &server_type { + Some(ConcreteDataType::Decimal128(dt)) => to_decimal_scalar_value(data, dt), + Some(st @ ConcreteDataType::Timestamp(unit)) => { + to_timestamp_scalar_value(data.and_then(|n| n.to_i64()), unit, st)? + } + Some(ConcreteDataType::UInt64(_)) | None => { + ScalarValue::UInt64(data.and_then(|n| n.to_u64())) + } + Some(st) => { + return Err(invalid_parameter_error( + "invalid_parameter_type", + Some(format!("Expected: {}, found: {}", st, client_type)), + )); + } + } + } &Type::FLOAT4 => { let data = portal.parameter::(idx, &client_type)?; if let Some(server_type) = &server_type { @@ -837,7 +811,67 @@ pub(super) fn parameters_to_scalar_values( ScalarValue::Null } } - &Type::VARCHAR_ARRAY => { + &Type::NUMERIC_ARRAY => { + let data = portal.parameter::>>(idx, &client_type)?; + if let Some(data) = data { + let build_u64_list = |data: Vec>| { + let values = data + .into_iter() + .map(|n| ScalarValue::UInt64(n.and_then(|n| n.to_u64()))) + .collect::>(); + ScalarValue::List(ScalarValue::new_list( + &values, + &ArrowDataType::UInt64, + true, + )) + }; + if let Some(server_type) = &server_type { + match server_type { + ConcreteDataType::List(list_type) => match list_type.item_type() { + ConcreteDataType::UInt64(_) => build_u64_list(data), + ConcreteDataType::Decimal128(dt) => { + let values = data + .into_iter() + .map(|n| to_decimal_scalar_value(n, dt)) + .collect::>(); + ScalarValue::List(ScalarValue::new_list( + &values, + &ArrowDataType::Decimal128(dt.precision(), dt.scale()), + true, + )) + } + _ => { + // the server type is not a list of decimal or uint64 + return Err(invalid_parameter_error( + "invalid_parameter_type", + Some(format!( + "Expected: {}, found: {}", + list_type.item_type(), + client_type + )), + )); + } + }, + _ => { + // the server type is not a list + return Err(invalid_parameter_error( + "invalid_parameter_type", + Some(format!( + "Expected: {}, found: {}", + server_type, client_type + )), + )); + } + } + } else { + // server type not provided + build_u64_list(data) + } + } else { + ScalarValue::Null + } + } + &Type::VARCHAR_ARRAY | &Type::TEXT_ARRAY | &Type::CHAR_ARRAY => { let data = portal.parameter::>>(idx, &client_type)?; if let Some(data) = data { let values = data.into_iter().map(|i| i.into()).collect::>(); @@ -1098,6 +1132,7 @@ pub fn format_options_from_query_ctx(query_ctx: &QueryContextRef) -> Arc Date: Fri, 10 Apr 2026 16:37:52 +0800 Subject: [PATCH 008/120] fix: fix current version comparison logic for pre-releases (#7946) Signed-off-by: liyang --- .github/scripts/check-version.sh | 69 +++++++++++++++++++++++++++++--- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/.github/scripts/check-version.sh b/.github/scripts/check-version.sh index 28c2812ded..1efa3bb4db 100755 --- a/.github/scripts/check-version.sh +++ b/.github/scripts/check-version.sh @@ -30,13 +30,72 @@ CLEAN_LATEST=$(echo "$LATEST_VERSION" | sed 's/^v//' | sed 's/-nightly-.*//') echo "Current version: $CLEAN_CURRENT" echo "Latest release version: $CLEAN_LATEST" -# Use sort -V to compare versions -HIGHER_VERSION=$(printf "%s\n%s" "$CLEAN_CURRENT" "$CLEAN_LATEST" | sort -V | tail -n1) +# Function to extract base version (without pre-release suffix) +get_base_version() { + echo "$1" | sed -E 's/-(alpha|beta|rc|pre).*//' +} -if [ "$HIGHER_VERSION" = "$CLEAN_CURRENT" ]; then +# Function to check if a version is pre-release +is_prerelease() { + [[ "$1" =~ -(alpha|beta|rc|pre) ]] +} + +# Compare versions properly considering pre-release +compare_versions() { + local current=$1 + local latest=$2 + + # Extract base versions + local current_base=$(get_base_version "$current") + local latest_base=$(get_base_version "$latest") + + # Compare base versions first + HIGHER_BASE=$(printf "%s\n%s" "$current_base" "$latest_base" | sort -V | tail -n1) + + if [ "$HIGHER_BASE" = "$latest_base" ] && [ "$current_base" != "$latest_base" ]; then + # Latest has higher base version + echo "current_older" + return + elif [ "$HIGHER_BASE" = "$current_base" ] && [ "$current_base" != "$latest_base" ]; then + # Current has higher base version + echo "current_newer" + return + fi + + # Base versions are equal, compare pre-release status + if [ "$current_base" = "$latest_base" ]; then + # If current is pre-release and latest is not, current is older + if is_prerelease "$current" && ! is_prerelease "$latest"; then + echo "current_older" + return + fi + + # If latest is pre-release and current is not, current is newer + if ! is_prerelease "$current" && is_prerelease "$latest"; then + echo "current_newer" + return + fi + fi + + # Both are same type or different base versions already handled, use sort -V + HIGHER_VERSION=$(printf "%s\n%s" "$current" "$latest" | sort -V | tail -n1) + if [ "$HIGHER_VERSION" = "$current" ]; then + echo "current_newer_or_equal" + else + echo "current_older" + fi +} + +RESULT=$(compare_versions "$CLEAN_CURRENT" "$CLEAN_LATEST") + +if [ "$RESULT" = "current_newer" ] || [ "$RESULT" = "current_newer_or_equal" ]; then echo "Current version ($CLEAN_CURRENT) is NEWER than or EQUAL to latest ($CLEAN_LATEST)" - echo "is-current-version-latest=true" >> $GITHUB_OUTPUT + if [ -n "$GITHUB_OUTPUT" ]; then + echo "is-current-version-latest=true" >> $GITHUB_OUTPUT + fi else echo "Current version ($CLEAN_CURRENT) is OLDER than latest ($CLEAN_LATEST)" - echo "is-current-version-latest=false" >> $GITHUB_OUTPUT + if [ -n "$GITHUB_OUTPUT" ]; then + echo "is-current-version-latest=false" >> $GITHUB_OUTPUT + fi fi From 06e49961c75a05a63e0c74e0fe8c909e25d6d134 Mon Sep 17 00:00:00 2001 From: cui Date: Fri, 10 Apr 2026 17:22:12 +0800 Subject: [PATCH 009/120] fix(index): intersect bitmaps before early exit in predicates applier (#7867) * fix(index): intersect bitmaps before early exit in predicates applier The loop skipped intersecting when the next bitmap was empty, which left the accumulator unchanged instead of zeroing it. Intersect first, then break when the result is empty. Signed-off-by: Weixie Cui * per gemini * style(index): format predicates applier loop * fix(index): remove unused mut in predicates applier --------- Signed-off-by: Weixie Cui Co-authored-by: discord9 <55937128+discord9@users.noreply.github.com> Co-authored-by: discord9 --- .../search/index_apply/predicates_apply.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index 441a4b4304..eaea8cfb95 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs @@ -94,15 +94,15 @@ impl IndexApplier for PredicatesIndexApplier { .collect::>(); let mut mapper = ParallelFstValuesMapper::new(reader); - let mut bm_vec = mapper.map_values_vec(&value_and_meta_vec, metrics).await?; + let bm_vec = mapper.map_values_vec(&value_and_meta_vec, metrics).await?; - let mut bitmap = bm_vec.pop().unwrap(); // SAFETY: `fst_ranges` is not empty - for bm in bm_vec { - if bm.count_ones() == 0 { + let mut iter = bm_vec.into_iter(); + let mut bitmap = iter.next().unwrap(); // SAFETY: `fst_ranges` is not empty + for bm in iter { + bitmap.intersect(bm); + if bitmap.count_ones() == 0 { break; } - - bitmap.intersect(bm); } output.matched_segment_ids = bitmap; From 76cad696c6ce26d9377491ca7ae46e078d6efbea Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Fri, 10 Apr 2026 18:41:48 +0800 Subject: [PATCH 010/120] feat: add parquet nested leaf projection (#7900) * feat: add parquet nested leaf projection * rename ParquetProjection related struct * add some apis * extract common build schema function for test * remove unsed method * keep only deduped parquet root projection constructor * add more unit tests * fix: typo * fix: cr * fast-path parquet root projection without nested fields * extract a build_projection_mask method * fix: cargo clippy --- src/mito2/src/sst/parquet.rs | 1 + src/mito2/src/sst/parquet/read_columns.rs | 316 ++++++++++++++++++++++ src/mito2/src/sst/parquet/reader.rs | 10 +- 3 files changed, 323 insertions(+), 4 deletions(-) create mode 100644 src/mito2/src/sst/parquet/read_columns.rs diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 2447824ad9..90395642b6 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -31,6 +31,7 @@ pub mod format; pub(crate) mod helper; pub(crate) mod metadata; pub mod prefilter; +pub mod read_columns; pub mod reader; pub mod row_group; pub mod row_selection; diff --git a/src/mito2/src/sst/parquet/read_columns.rs b/src/mito2/src/sst/parquet/read_columns.rs new file mode 100644 index 0000000000..f0f35a4099 --- /dev/null +++ b/src/mito2/src/sst/parquet/read_columns.rs @@ -0,0 +1,316 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; + +/// A nested field access path inside one parquet root column. +pub type ParquetNestedPath = Vec; + +/// The parquet columns to read. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParquetReadColumns { + cols: Vec, + has_nested: bool, +} + +impl ParquetReadColumns { + /// Builds root-column projections from root indices that are already + /// deduplicated. + /// + /// Note: this constructor does not check for duplicates. + pub fn from_deduped_root_indices(root_indices: impl IntoIterator) -> Self { + let cols = root_indices + .into_iter() + .map(ParquetReadColumn::new) + .collect(); + Self { + cols, + has_nested: false, + } + } + + pub fn columns(&self) -> &[ParquetReadColumn] { + &self.cols + } + + pub fn has_nested(&self) -> bool { + self.has_nested + } + + pub fn root_indices_iter(&self) -> impl Iterator + '_ { + self.cols.iter().map(|col| col.root_index) + } +} + +/// Read requirement for a single parquet root column. +/// +/// `root_index` identifies the root column in the parquet schema. +/// +/// If `nested_paths` is empty, the whole root column is read. Otherwise, only +/// leaves under the specified nested paths are read. +/// +/// To construct a [`ParquetReadColumn`]: +/// - `ParquetReadColumn::new(0)` reads the whole root column at index `0`. +/// - `ParquetReadColumn::new(0).with_nested_paths(vec![vec!["j".into(), "b".into()]])` +/// reads only leaves under `j.b`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParquetReadColumn { + /// Root field index in the parquet schema. + root_index: usize, + /// Nested paths to read under this root column. + /// + /// Each path includes the root column itself. For example, for a root + /// column `j`, path `["j", "a", "b"]` refers to `j.a.b`. + /// + /// If empty, the whole root column is read. + nested_paths: Vec, +} + +impl ParquetReadColumn { + pub fn new(root_index: usize) -> Self { + Self { + root_index, + nested_paths: vec![], + } + } + + pub fn with_nested_paths(self, nested_paths: Vec) -> Self { + Self { + nested_paths, + ..self + } + } + + pub fn root_index(&self) -> usize { + self.root_index + } + + pub fn nested_paths(&self) -> &[ParquetNestedPath] { + &self.nested_paths + } +} + +/// Builds a projection mask from parquet read columns. +pub fn build_projection_mask( + parquet_read_cols: &ParquetReadColumns, + parquet_schema_desc: &SchemaDescriptor, +) -> ProjectionMask { + if parquet_read_cols.has_nested() { + let leaf_indices = build_parquet_leaves_indices(parquet_schema_desc, parquet_read_cols); + ProjectionMask::leaves(parquet_schema_desc, leaf_indices) + } else { + ProjectionMask::roots(parquet_schema_desc, parquet_read_cols.root_indices_iter()) + } +} + +/// Builds parquet leaf-column indices from parquet read columns. +fn build_parquet_leaves_indices( + parquet_schema_desc: &SchemaDescriptor, + projection: &ParquetReadColumns, +) -> Vec { + let mut map = HashMap::with_capacity(projection.cols.len()); + for col in &projection.cols { + map.insert(col.root_index, &col.nested_paths); + } + + let mut leaf_indices = Vec::new(); + for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() { + let root_idx = parquet_schema_desc.get_column_root_idx(leaf_idx); + let Some(nested_paths) = map.get(&root_idx) else { + continue; + }; + if nested_paths.is_empty() { + leaf_indices.push(leaf_idx); + continue; + } + + let leaf_path = leaf_col.path().parts(); + if nested_paths + .iter() + .any(|nested_path| leaf_path.starts_with(nested_path)) + { + leaf_indices.push(leaf_idx); + } + } + leaf_indices +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use parquet::basic::Repetition; + use parquet::schema::types::Type; + + use super::*; + + #[test] + fn test_reads_whole_root() { + let parquet_schema_desc = build_test_nested_parquet_schema(); + + let projection = ParquetReadColumns { + cols: vec![ParquetReadColumn { + root_index: 0, + nested_paths: vec![], + }], + has_nested: false, + }; + + assert_eq!( + vec![0, 1, 2], + build_parquet_leaves_indices(&parquet_schema_desc, &projection) + ); + } + + #[test] + fn test_filters_nested_paths() { + let parquet_schema_desc = build_test_nested_parquet_schema(); + + let projection = ParquetReadColumns { + cols: vec![ + ParquetReadColumn { + root_index: 0, + nested_paths: vec![vec!["j".to_string(), "b".to_string()]], + }, + ParquetReadColumn { + root_index: 1, + nested_paths: vec![], + }, + ], + has_nested: true, + }; + + assert_eq!( + vec![1, 2, 3], + build_parquet_leaves_indices(&parquet_schema_desc, &projection) + ); + } + + #[test] + fn test_reads_middle_level_path() { + let parquet_schema_desc = build_test_nested_parquet_schema(); + + let projection = ParquetReadColumns { + cols: vec![ParquetReadColumn { + root_index: 0, + nested_paths: vec![vec!["j".to_string(), "b".to_string()]], + }], + has_nested: true, + }; + + assert_eq!( + vec![1, 2], + build_parquet_leaves_indices(&parquet_schema_desc, &projection) + ); + } + + #[test] + fn test_reads_leaf_level_path() { + let parquet_schema_desc = build_test_nested_parquet_schema(); + + let projection = ParquetReadColumns { + cols: vec![ParquetReadColumn { + root_index: 0, + nested_paths: vec![vec!["j".to_string(), "b".to_string(), "c".to_string()]], + }], + has_nested: true, + }; + + assert_eq!( + vec![1], + build_parquet_leaves_indices(&parquet_schema_desc, &projection) + ); + } + + #[test] + fn test_merges_mixed_paths() { + let parquet_schema_desc = build_test_nested_parquet_schema(); + + let projection = ParquetReadColumns { + cols: vec![ParquetReadColumn { + root_index: 0, + nested_paths: vec![ + vec!["j".to_string(), "a".to_string()], + vec!["j".to_string(), "b".to_string(), "d".to_string()], + ], + }], + has_nested: true, + }; + + assert_eq!( + vec![0, 2], + build_parquet_leaves_indices(&parquet_schema_desc, &projection) + ); + } + + // Test schema: + // schema + // |- j + // | |- a: INT64 + // | `- b + // | |- c: INT64 + // | `- d: INT64 + // `- k: INT64 + fn build_test_nested_parquet_schema() -> SchemaDescriptor { + let leaf_a = Arc::new( + Type::primitive_type_builder("a", parquet::basic::Type::INT64) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ); + let leaf_c = Arc::new( + Type::primitive_type_builder("c", parquet::basic::Type::INT64) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ); + let leaf_d = Arc::new( + Type::primitive_type_builder("d", parquet::basic::Type::INT64) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ); + let group_b = Arc::new( + Type::group_type_builder("b") + .with_repetition(Repetition::REQUIRED) + .with_fields(vec![leaf_c, leaf_d]) + .build() + .unwrap(), + ); + let root_j = Arc::new( + Type::group_type_builder("j") + .with_repetition(Repetition::REQUIRED) + .with_fields(vec![leaf_a, group_b]) + .build() + .unwrap(), + ); + let root_k = Arc::new( + Type::primitive_type_builder("k", parquet::basic::Type::INT64) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ); + let schema = Arc::new( + Type::group_type_builder("schema") + .with_fields(vec![root_j, root_k]) + .build() + .unwrap(), + ); + + SchemaDescriptor::new(schema) + } +} diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 73ca7748e9..6fdbb6f243 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -79,6 +79,7 @@ use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{ PrefilterContextBuilder, execute_prefilter, is_usable_primary_key_filter, }; +use crate::sst::parquet::read_columns::{ParquetReadColumns, build_projection_mask}; use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::sst::parquet::row_selection::RowGroupSelection; use crate::sst::parquet::stats::RowGroupPruningStats; @@ -406,10 +407,11 @@ impl ParquetReaderBuilder { // Computes the projection mask. let parquet_schema_desc = parquet_meta.file_metadata().schema_descr(); - let indices = read_format.projection_indices(); - // Now we assumes we don't have nested schemas. - // TODO(yingwen): Revisit this if we introduce nested types such as JSON type. - let projection_mask = ProjectionMask::roots(parquet_schema_desc, indices.iter().copied()); + let parquet_read_cols = ParquetReadColumns::from_deduped_root_indices( + read_format.projection_indices().iter().copied(), + ); + + let projection_mask = build_projection_mask(&parquet_read_cols, parquet_schema_desc); let selection = self .row_groups_to_read(&read_format, &parquet_meta, &mut metrics.filter_metrics) .await; From d1b2a310975d1c4d54eac746563ca87b74b0b9d1 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 13 Apr 2026 14:42:55 +0800 Subject: [PATCH 011/120] fix: randomize standalone test ports in cli export test (#7955) fix/flaky-test: ### Add Dynamic Port Selection for Standalone Tests - **`cli.rs`**: Implemented functions `random_standalone_addrs` and `choose_random_unused_port_offset` to dynamically select unused ports for standalone tests, enhancing test reliability. - Updated `test_export_create_table_with_quoted_names` to use dynamically assigned ports for HTTP, RPC, MySQL, and PostgreSQL addresses. Signed-off-by: Lei, HUANG --- src/cmd/src/cli.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/src/cmd/src/cli.rs b/src/cmd/src/cli.rs index 84e797c291..95c5f00b77 100644 --- a/src/cmd/src/cli.rs +++ b/src/cmd/src/cli.rs @@ -102,31 +102,79 @@ impl Command { #[cfg(test)] mod tests { + use std::net::TcpListener; + use std::ops::RangeInclusive; + use clap::Parser; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_telemetry::logging::LoggingOptions; + use rand::Rng; use crate::error::Result as CmdResult; use crate::options::GlobalOptions; use crate::{App, cli, standalone}; + fn random_standalone_addrs() -> (String, String, String, String) { + let offset = choose_random_unused_port_offset(14000..=24000, 10); + + ( + format!("127.0.0.1:{}", 4000 + offset), + format!("127.0.0.1:{}", 4001 + offset), + format!("127.0.0.1:{}", 4002 + offset), + format!("127.0.0.1:{}", 4003 + offset), + ) + } + + fn choose_random_unused_port_offset( + port_range: RangeInclusive, + max_attempts: usize, + ) -> u16 { + let mut rng = rand::rng(); + + for _ in 0..max_attempts { + let http_port = rng.random_range(port_range.clone()); + let offset = http_port - 4000; + let ports = [4000 + offset, 4001 + offset, 4002 + offset, 4003 + offset]; + + let listeners = ports + .into_iter() + .map(|port| TcpListener::bind(("127.0.0.1", port))) + .collect::, _>>(); + + if listeners.is_ok() { + return offset; + } + } + + panic!("failed to find unused standalone test ports"); + } + #[tokio::test(flavor = "multi_thread")] async fn test_export_create_table_with_quoted_names() -> CmdResult<()> { let output_dir = tempfile::tempdir().unwrap(); + let (http_addr, rpc_addr, mysql_addr, postgres_addr) = random_standalone_addrs(); let standalone = standalone::Command::parse_from([ "standalone", "start", "--data-home", &*output_dir.path().to_string_lossy(), + "--http-addr", + &http_addr, + "--rpc-bind-addr", + &rpc_addr, + "--mysql-addr", + &mysql_addr, + "--postgres-addr", + &postgres_addr, ]); let standalone_opts = standalone.load_options(&GlobalOptions::default()).unwrap(); let mut instance = standalone.build(standalone_opts).await?; instance.start().await?; - let client = Client::with_urls(["127.0.0.1:4001"]); + let client = Client::with_urls([rpc_addr.as_str()]); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); database .sql(r#"CREATE DATABASE "cli.export.create_table";"#) @@ -149,7 +197,7 @@ mod tests { "data", "export", "--addr", - "127.0.0.1:4000", + &http_addr, "--output-dir", &*output_dir.path().to_string_lossy(), "--target", From 9f7ffb4d26f498ceb8ea24349f6e341fe29f8365 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 13 Apr 2026 16:12:11 +0800 Subject: [PATCH 012/120] feat(mito2): allow CompactionOutput to succeed independently (#7948) * refactor(mito2): improve compaction error handling and file removal Refactor compaction task execution to enhance error handling and robustness. - Implemented parallel execution of compaction tasks with proper error capture and logging for individual task failures. - Ensured JoinSnafu is no longer directly used in error propagation, instead handling errors within the task processing loop. - Adjusted file removal logic to correctly include expired SSTs after compaction merges. Signed-off-by: Lei, HUANG * refactor(mito2): extract SstMerger trait for testability in compaction Extract SstMerger trait and DefaultSstMerger implementation to improve the testability of DefaultCompactor. The DefaultCompactor is now generic over SstMerger, allowing mock implementations to be injected for unit testing without relying on the full object storage access layer. This refactoring separates the concerns of SST file merging from the overall compaction orchestration logic. Additionally: - Updated CompactionScheduler to use DefaultCompactor::default(). - Added unit tests for DefaultCompactor using a MockMerger. Signed-off-by: Lei, HUANG * fix(compaction): propagate join error during sst flush Correctly propagates the error when joining SST flush handles during compaction. Previously, the error was logged but not returned, leading to potential silent failures. Also reorders some imports for consistency. Signed-off-by: Lei, HUANG * perf(compaction): pre-allocate capacity for compacted_inputs Pre-allocates capacity for the compacted_inputs vector based on the estimated total size of inputs and expired SSTs. This optimization aims to reduce vector reallocations during the compaction process. Signed-off-by: Lei, HUANG * feat/allow-partial-compaction: ### Commit Message Enhance `DefaultCompactor` and `MockMerger` for Improved Flexibility - **`compactor.rs`**: - Added `Clone` trait to `DefaultSstMerger` and `MockMerger` to allow cloning. - Removed `Arc` wrapping from `DefaultCompactor`'s `merger` field for direct usage. - Updated `merge_ssts` method to require `Clone` trait for `SstMerger`. - Modified `MockMerger` to use `Arc` for `results` and `call_idx` to ensure thread safety. - Adjusted error handling to use `error::InvalidMetaSnafu` directly. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/compaction.rs | 2 +- src/mito2/src/compaction/compactor.rs | 390 +++++++++++++++++++++++--- 2 files changed, 354 insertions(+), 38 deletions(-) diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 944c51ebd6..d2120690ac 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -562,7 +562,7 @@ impl CompactionScheduler { listener, picker_output, compaction_region, - compactor: Arc::new(DefaultCompactor {}), + compactor: Arc::new(DefaultCompactor::default()), memory_manager: self.memory_manager.clone(), memory_policy: self.memory_policy, estimated_memory_bytes: estimated_bytes, diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index ff4317331f..fd3d01b276 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -38,11 +38,10 @@ use crate::compaction::picker::{PickerOutput, new_picker}; use crate::compaction::{CompactionOutput, CompactionSstReaderBuilder, find_dynamic_options}; use crate::config::MitoConfig; use crate::error::{ - EmptyRegionDirSnafu, InvalidPartitionExprSnafu, JoinSnafu, ObjectStoreNotFoundSnafu, Result, + EmptyRegionDirSnafu, InvalidPartitionExprSnafu, ObjectStoreNotFoundSnafu, Result, }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; -use crate::metrics; use crate::read::FlatSource; use crate::region::options::RegionOptions; use crate::region::version::VersionRef; @@ -56,6 +55,7 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::sst::location::region_dir_from_table_dir; use crate::sst::parquet::WriteOptions; use crate::sst::version::{SstVersion, SstVersionRef}; +use crate::{error, metrics}; /// Region version for compaction that does not hold memtables. #[derive(Clone)] @@ -299,12 +299,28 @@ pub trait Compactor: Send + Sync + 'static { ) -> Result<()>; } -/// DefaultCompactor is the default implementation of Compactor. -pub struct DefaultCompactor; - -impl DefaultCompactor { - /// Merge a single compaction output into SST files. +/// Trait for merging a single compaction output into SST files. +/// +/// This is extracted from `DefaultCompactor` to allow injecting mock +/// implementations in tests. +#[async_trait::async_trait] +pub trait SstMerger: Send + Sync + 'static { async fn merge_single_output( + &self, + compaction_region: CompactionRegion, + output: CompactionOutput, + write_opts: WriteOptions, + ) -> Result>; +} + +/// The production [`SstMerger`] that reads, merges, and writes SST files. +#[derive(Clone)] +pub struct DefaultSstMerger; + +#[async_trait::async_trait] +impl SstMerger for DefaultSstMerger { + async fn merge_single_output( + &self, compaction_region: CompactionRegion, output: CompactionOutput, write_opts: WriteOptions, @@ -424,54 +440,113 @@ impl DefaultCompactor { } } +/// DefaultCompactor is the default implementation of Compactor. +/// +/// It is parameterized by an [`SstMerger`] to allow injecting mock +/// implementations in tests. +pub struct DefaultCompactor { + merger: M, +} + +impl Default for DefaultCompactor { + fn default() -> Self { + Self { + merger: DefaultSstMerger, + } + } +} + +impl DefaultCompactor { + pub fn with_merger(merger: M) -> Self { + Self { merger } + } +} + #[async_trait::async_trait] -impl Compactor for DefaultCompactor { +impl Compactor for DefaultCompactor +where + M: Clone, +{ async fn merge_ssts( &self, compaction_region: &CompactionRegion, mut picker_output: PickerOutput, ) -> Result { - let mut futs = Vec::with_capacity(picker_output.outputs.len()); - let mut compacted_inputs = - Vec::with_capacity(picker_output.outputs.iter().map(|o| o.inputs.len()).sum()); let internal_parallelism = compaction_region.max_parallelism.max(1); let compaction_time_window = picker_output.time_window_size; + let region_id = compaction_region.region_id; + + // Build tasks along with their input file metas so we can track which + // inputs correspond to each task. + let mut tasks: Vec<(Vec, _)> = Vec::with_capacity(picker_output.outputs.len()); for output in picker_output.outputs.drain(..) { let inputs_to_remove: Vec<_> = output.inputs.iter().map(|f| f.meta_ref().clone()).collect(); - compacted_inputs.extend(inputs_to_remove.iter().cloned()); let write_opts = WriteOptions { write_buffer_size: compaction_region.engine_config.sst_write_buffer_size, max_file_size: picker_output.max_file_size, ..Default::default() }; - futs.push(Self::merge_single_output( - compaction_region.clone(), - output, - write_opts, - )); - } - let mut output_files = Vec::with_capacity(futs.len()); - while !futs.is_empty() { - let mut task_chunk = Vec::with_capacity(internal_parallelism); - for _ in 0..internal_parallelism { - if let Some(task) = futs.pop() { - task_chunk.push(common_runtime::spawn_compact(task)); - } - } - let metas = futures::future::try_join_all(task_chunk) - .await - .context(JoinSnafu)? - .into_iter() - .collect::>>>()?; - output_files.extend(metas.into_iter().flatten()); + let merger = self.merger.clone(); + let compaction_region = compaction_region.clone(); + let fut = async move { + merger + .merge_single_output(compaction_region, output, write_opts) + .await + }; + tasks.push((inputs_to_remove, fut)); } - // In case of remote compaction, we still allow the region edit after merge to - // clean expired ssts. - let mut inputs: Vec<_> = compacted_inputs.into_iter().collect(); - inputs.extend( + let mut output_files = Vec::with_capacity(tasks.len()); + let mut compacted_inputs = Vec::with_capacity( + tasks.iter().map(|(inputs, _)| inputs.len()).sum::() + + picker_output.expired_ssts.len(), + ); + + while !tasks.is_empty() { + let mut chunk: Vec<(Vec, _)> = Vec::with_capacity(internal_parallelism); + for _ in 0..internal_parallelism { + if let Some(task) = tasks.pop() { + chunk.push(task); + } + } + let spawned: Vec<_> = chunk + .into_iter() + .map(|(inputs, fut)| { + let handle = common_runtime::spawn_compact(fut); + (inputs, handle) + }) + .collect(); + + for (inputs, handle) in spawned { + match handle.await { + Ok(Ok(files)) => { + output_files.extend(files); + compacted_inputs.extend(inputs); + } + Ok(Err(e)) => { + warn!( + e; "Region {} failed to merge compaction output with inputs: [{}], skipping", + region_id, + inputs.iter().map(|f| f.file_id.to_string()).join(",") + ); + } + Err(e) => { + warn!( + "Region {} compaction task join error for inputs: [{}], skipping: {}", + region_id, + inputs.iter().map(|f| f.file_id.to_string()).join(","), + e + ); + return Err(e).context(error::JoinSnafu); + } + } + } + } + + // Include expired SSTs in removals — these don't depend on merge success. + compacted_inputs.extend( picker_output .expired_ssts .iter() @@ -480,7 +555,7 @@ impl Compactor for DefaultCompactor { Ok(MergeOutput { files_to_add: output_files, - files_to_remove: inputs, + files_to_remove: compacted_inputs, compaction_time_window: Some(compaction_time_window), }) } @@ -558,3 +633,244 @@ impl Compactor for DefaultCompactor { Ok(()) } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use store_api::storage::{FileId, RegionId}; + + use super::*; + use crate::cache::CacheManager; + use crate::compaction::picker::PickerOutput; + use crate::sst::file::FileHandle; + use crate::sst::file_purger::NoopFilePurger; + use crate::sst::version::SstVersion; + use crate::test_util::memtable_util::metadata_for_test; + use crate::test_util::scheduler_util::SchedulerEnv; + + fn dummy_file_meta() -> FileMeta { + FileMeta { + region_id: RegionId::new(1, 1), + file_id: FileId::random(), + file_size: 100, + ..Default::default() + } + } + + fn new_file_handle(meta: FileMeta) -> FileHandle { + FileHandle::new(meta, Arc::new(NoopFilePurger)) + } + + /// Build a minimal [`CompactionRegion`] suitable for tests where the + /// [`SstMerger`] is mocked and never touches the access layer. + async fn new_test_compaction_region() -> CompactionRegion { + let env = SchedulerEnv::new().await; + let metadata = metadata_for_test(); + let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; + CompactionRegion { + region_id: RegionId::new(1, 1), + region_options: RegionOptions::default(), + engine_config: Arc::new(MitoConfig::default()), + region_metadata: metadata.clone(), + cache_manager: Arc::new(CacheManager::default()), + access_layer: env.access_layer.clone(), + manifest_ctx, + current_version: CompactionVersion { + metadata, + options: RegionOptions::default(), + ssts: Arc::new(SstVersion::new()), + compaction_time_window: None, + }, + file_purger: None, + ttl: None, + max_parallelism: 1, + } + } + + /// An [`SstMerger`] that returns pre-configured results per call index. + /// + /// Call 0 gets `results[0]`, call 1 gets `results[1]`, etc. + #[derive(Clone)] + struct MockMerger { + results: Arc>>>>, + call_idx: Arc, + } + + impl MockMerger { + fn new(results: Vec>>) -> Self { + Self { + results: Arc::new(Mutex::new(results)), + call_idx: Arc::new(AtomicUsize::new(0)), + } + } + } + + #[async_trait::async_trait] + impl SstMerger for MockMerger { + async fn merge_single_output( + &self, + _compaction_region: CompactionRegion, + _output: CompactionOutput, + _write_opts: WriteOptions, + ) -> Result> { + let idx = self.call_idx.fetch_add(1, Ordering::SeqCst); + match self.results.lock().unwrap().get(idx) { + Some(Ok(files)) => Ok(files.clone()), + Some(Err(_)) => error::InvalidMetaSnafu { + reason: format!("simulated failure at index {idx}"), + } + .fail(), + None => panic!("MockMerger: no result configured for call index {idx}"), + } + } + } + + #[tokio::test] + async fn test_partial_merge_failure_collects_only_successful_outputs() { + common_telemetry::init_default_ut_logging(); + + let compaction_region = new_test_compaction_region().await; + + // Prepare 3 compaction outputs: output 0 and 2 succeed, output 1 fails. + let input_meta_0 = dummy_file_meta(); + let input_meta_1 = dummy_file_meta(); + let input_meta_2 = dummy_file_meta(); + + let output_meta_0 = vec![dummy_file_meta()]; + let output_meta_2 = vec![dummy_file_meta(), dummy_file_meta()]; + + let merger = MockMerger::new(vec![ + Ok(output_meta_0.clone()), + Err(error::InvalidMetaSnafu { + reason: "boom".to_string(), + } + .build()), + Ok(output_meta_2.clone()), + ]); + let compactor = DefaultCompactor::with_merger(merger); + + let picker_output = PickerOutput { + outputs: vec![ + CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input_meta_0.clone())], + filter_deleted: false, + output_time_range: None, + }, + CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input_meta_1.clone())], + filter_deleted: false, + output_time_range: None, + }, + CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input_meta_2.clone())], + filter_deleted: false, + output_time_range: None, + }, + ], + expired_ssts: vec![], + time_window_size: 3600, + max_file_size: None, + }; + + let merge_output = compactor + .merge_ssts(&compaction_region, picker_output) + .await + .unwrap(); + + // Outputs 0 and 2 succeeded (1 + 2 = 3 files added). + assert_eq!(merge_output.files_to_add.len(), 3); + // Only inputs from successful merges should be removed. + assert_eq!(merge_output.files_to_remove.len(), 2); + + let removed_ids: Vec<_> = merge_output + .files_to_remove + .iter() + .map(|f| f.file_id) + .collect(); + assert!(removed_ids.contains(&input_meta_0.file_id)); + assert!(removed_ids.contains(&input_meta_2.file_id)); + // The failed output's input must NOT be removed. + assert!(!removed_ids.contains(&input_meta_1.file_id)); + } + + #[tokio::test] + async fn test_all_outputs_succeed() { + common_telemetry::init_default_ut_logging(); + + let compaction_region = new_test_compaction_region().await; + let input_meta = dummy_file_meta(); + let output_meta = vec![dummy_file_meta()]; + + let merger = MockMerger::new(vec![Ok(output_meta.clone())]); + let compactor = DefaultCompactor::with_merger(merger); + + let picker_output = PickerOutput { + outputs: vec![CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input_meta.clone())], + filter_deleted: false, + output_time_range: None, + }], + expired_ssts: vec![], + time_window_size: 3600, + max_file_size: None, + }; + + let merge_output = compactor + .merge_ssts(&compaction_region, picker_output) + .await + .unwrap(); + + assert_eq!(merge_output.files_to_add.len(), 1); + assert_eq!(merge_output.files_to_add[0].file_id, output_meta[0].file_id); + assert_eq!(merge_output.files_to_remove.len(), 1); + assert_eq!(merge_output.files_to_remove[0].file_id, input_meta.file_id); + } + + #[tokio::test] + async fn test_expired_ssts_always_removed() { + common_telemetry::init_default_ut_logging(); + + let compaction_region = new_test_compaction_region().await; + let input_meta = dummy_file_meta(); + let expired_meta = dummy_file_meta(); + + // The single merge output fails, but expired SSTs should still be removed. + let merger = MockMerger::new(vec![Err(error::InvalidMetaSnafu { + reason: "fail".to_string(), + } + .build())]); + let compactor = DefaultCompactor::with_merger(merger); + + let picker_output = PickerOutput { + outputs: vec![CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input_meta.clone())], + filter_deleted: false, + output_time_range: None, + }], + expired_ssts: vec![new_file_handle(expired_meta.clone())], + time_window_size: 3600, + max_file_size: None, + }; + + let merge_output = compactor + .merge_ssts(&compaction_region, picker_output) + .await + .unwrap(); + + // No files added (merge failed). + assert!(merge_output.files_to_add.is_empty()); + // Only the expired SST should be in files_to_remove (not the failed merge's input). + assert_eq!(merge_output.files_to_remove.len(), 1); + assert_eq!( + merge_output.files_to_remove[0].file_id, + expired_meta.file_id + ); + } +} From 01a73105b8e5441ee66170cf3816085f68d9019e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 13 Apr 2026 16:27:53 +0800 Subject: [PATCH 013/120] feat: use partition range cache in scan (#7873) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: use range cache in scan Signed-off-by: evenyag * refactor: rename dedup to skip_dedup Signed-off-by: evenyag * feat: use background concat for buffered batches Signed-off-by: evenyag * chore: fmt Signed-off-by: evenyag * fix: store permits Signed-off-by: evenyag * fix: fix potential panic Signed-off-by: evenyag * fix: skip range-cache wrapping when cache is disabled Signed-off-by: evenyag * fix: avoid potential deadlock Deadlock Chain 1. Range-level merge tasks: Each concurrent build_flat_partition_range_read (line 494-506) calls build_flat_reader_from_sources → create_parallel_flat_sources → spawn_flat_scan_task. These background tasks loop: acquire permit → input.next() → release permit. 2. Final merge tasks: After all range tasks return streams (line 509-511), the distributor calls build_flat_reader_from_sources again (line 520-527) → create_parallel_flat_sources → more spawn_flat_scan_task tasks. These also loop: acquire permit → input.next() → release permit. 3. Circular wait: The final merge tasks' input.next() reads from ReceiverStreams backed by range-level merge tasks. If all num_partitions permits are held by final merge tasks blocked on input.next(), the range-level merge tasks can't acquire permits to produce data → deadlock. Signed-off-by: evenyag * test: add test for small permits Signed-off-by: evenyag * feat: use avg batch size for channel size Signed-off-by: evenyag * test: fix test Signed-off-by: evenyag * chore: address review comments Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/cache.rs | 70 ++++ src/mito2/src/engine/scan_test.rs | 96 +++++ src/mito2/src/read/range_cache.rs | 658 +++++++++++++++++++----------- src/mito2/src/read/scan_util.rs | 139 ++++++- src/mito2/src/read/seq_scan.rs | 83 +++- src/mito2/src/read/series_scan.rs | 85 ++-- 6 files changed, 845 insertions(+), 286 deletions(-) diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 35db74eee6..5d2559cba1 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -28,6 +28,7 @@ use std::ops::Range; use std::sync::Arc; use bytes::Bytes; +use common_base::readable_size::ReadableSize; use common_telemetry::warn; use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; @@ -72,6 +73,46 @@ const INDEX_TYPE: &str = "index"; const SELECTOR_RESULT_TYPE: &str = "selector_result"; /// Metrics type key for range scan result cache. const RANGE_RESULT_TYPE: &str = "range_result"; +const RANGE_RESULT_CONCAT_MEMORY_LIMIT: ReadableSize = ReadableSize::mb(512); +const RANGE_RESULT_CONCAT_MEMORY_PERMIT: ReadableSize = ReadableSize::kb(1); + +#[derive(Debug)] +pub(crate) struct RangeResultMemoryLimiter { + semaphore: Arc, + permit_bytes: usize, +} + +impl Default for RangeResultMemoryLimiter { + fn default() -> Self { + Self::new( + RANGE_RESULT_CONCAT_MEMORY_LIMIT.as_bytes() as usize, + RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize, + ) + } +} + +impl RangeResultMemoryLimiter { + pub(crate) fn new(limit_bytes: usize, permit_bytes: usize) -> Self { + let permit_bytes = permit_bytes.max(1); + let permits = limit_bytes.div_ceil(permit_bytes).max(1); + Self { + semaphore: Arc::new(tokio::sync::Semaphore::new(permits)), + permit_bytes, + } + } + + pub(crate) fn permit_bytes(&self) -> usize { + self.permit_bytes + } + + pub(crate) async fn acquire( + &self, + bytes: usize, + ) -> std::result::Result, tokio::sync::AcquireError> { + let permits = bytes.div_ceil(self.permit_bytes()).max(1) as u32; + self.semaphore.acquire_many(permits).await + } +} /// Cached SST metadata combines the parquet footer with the decoded region metadata. /// @@ -373,6 +414,23 @@ impl CacheStrategy { } } + /// Returns true if the range result cache is enabled. + pub(crate) fn has_range_result_cache(&self) -> bool { + match self { + CacheStrategy::EnableAll(cache_manager) => cache_manager.has_range_result_cache(), + CacheStrategy::Compaction(_) | CacheStrategy::Disabled => false, + } + } + + pub(crate) fn range_result_memory_limiter(&self) -> Option<&Arc> { + match self { + CacheStrategy::EnableAll(cache_manager) => { + Some(cache_manager.range_result_memory_limiter()) + } + CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, + } + } + /// Calls [CacheManager::write_cache()]. /// It returns None if the strategy is [CacheStrategy::Disabled]. pub fn write_cache(&self) -> Option<&WriteCacheRef> { @@ -476,6 +534,8 @@ pub struct CacheManager { selector_result_cache: Option, /// Cache for range scan outputs in flat format. range_result_cache: Option, + /// Shared memory limiter for async range-result cache tasks. + range_result_memory_limiter: Arc, /// Cache for index result. index_result_cache: Option, } @@ -735,6 +795,15 @@ impl CacheManager { } } + /// Returns true if the range result cache is enabled. + pub(crate) fn has_range_result_cache(&self) -> bool { + self.range_result_cache.is_some() + } + + pub(crate) fn range_result_memory_limiter(&self) -> &Arc { + &self.range_result_memory_limiter + } + /// Gets the write cache. pub(crate) fn write_cache(&self) -> Option<&WriteCacheRef> { self.write_cache.as_ref() @@ -969,6 +1038,7 @@ impl CacheManagerBuilder { puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)), selector_result_cache, range_result_cache, + range_result_memory_limiter: Arc::new(RangeResultMemoryLimiter::default()), index_result_cache, } } diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 119b4493fd..a39761ad01 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -403,3 +403,99 @@ fn collect_and_assert_partition_rows( actual_rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2))); actual_rows } + +/// Tests series scan with multiple partition ranges (each with multiple overlapping sources) +/// and small semaphore permits (controlled by num_partitions). +#[tokio::test] +async fn test_series_scan_flat_small_permits() { + let mut env = TestEnv::with_prefix("test_series_scan_small_permits").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .build(); + let column_schemas = test_util::rows_schema(&request); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Create overlapping SSTs in each time window so partition ranges have multiple sources. + let put_flush_rows = async |start, end| { + let rows = Rows { + schema: column_schemas.clone(), + rows: test_util::build_rows(start, end), + }; + test_util::put_rows(&engine, region_id, rows).await; + test_util::flush_region(&engine, region_id, None).await; + }; + // Window 0 (0s-999s): 3 overlapping SSTs + put_flush_rows(0, 3).await; + put_flush_rows(1, 5).await; + put_flush_rows(3, 7).await; + // Window 1 (3600s-4599s): 2 overlapping SSTs + put_flush_rows(3600, 3603).await; + put_flush_rows(3601, 3605).await; + // Window 2 (7200s-8199s): 2 overlapping SSTs + put_flush_rows(7200, 7203).await; + put_flush_rows(7201, 7204).await; + + let mut expected_rows = Vec::new(); + for value in [ + 0_i64, 1, 2, 3, 4, 5, 6, 3600, 3601, 3602, 3603, 3604, 7200, 7201, 7202, 7203, + ] { + expected_rows.push((value.to_string(), value as f64, value * 1000)); + } + expected_rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2))); + + // Test with different semaphore sizes (num_partitions controls Semaphore::new(num_partitions)). + for num_partitions in [1, 2] { + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await.unwrap(); + let Scanner::Series(mut scanner) = scanner else { + panic!("Scanner should be series scan"); + }; + + // Collect all partition ranges and redistribute into `num_partitions` partitions. + let raw_ranges: Vec<_> = scanner + .properties() + .partitions + .iter() + .flatten() + .cloned() + .collect(); + assert!( + raw_ranges.len() >= 3, + "expected at least 3 partition ranges, got {}", + raw_ranges.len() + ); + + let mut new_ranges = vec![vec![]; num_partitions]; + for (i, range) in raw_ranges.into_iter().enumerate() { + new_ranges[i % num_partitions].push(range); + } + scanner + .prepare(PrepareRequest { + ranges: Some(new_ranges), + ..Default::default() + }) + .unwrap(); + + let actual_rows = collect_partition_rows_round_robin(&scanner, num_partitions).await; + assert_eq!( + expected_rows, actual_rows, + "mismatch with num_partitions={num_partitions}" + ); + } +} diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 2431a21f6a..1daaa6399b 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -18,22 +18,28 @@ use std::mem; use std::sync::Arc; use async_stream::try_stream; +use common_telemetry::warn; use common_time::range::TimestampRange; -use datatypes::arrow::array::{Array, AsArray, DictionaryArray}; -use datatypes::arrow::datatypes::UInt32Type; +use datatypes::arrow::compute::concat_batches; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use futures::TryStreamExt; +use snafu::ResultExt; use store_api::region_engine::PartitionRange; use store_api::storage::{ColumnId, FileId, RegionId, TimeSeriesRowSelector}; +use tokio::sync::{mpsc, oneshot}; use crate::cache::CacheStrategy; +use crate::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use crate::read::BoxedRecordBatchStream; use crate::read::scan_region::StreamContext; use crate::read::scan_util::PartitionMetrics; use crate::region::options::MergeMode; use crate::sst::file::FileTimeRange; -use crate::sst::parquet::flat_format::primary_key_column_index; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; + +const RANGE_CACHE_COMPACT_THRESHOLD_BYTES: usize = 2 * 1024 * 1024; +const RANGE_CACHE_SKIP_BYTES: usize = 512 * 1024 * 1024; /// Fingerprint of the scan request fields that affect partition range cache reuse. /// @@ -187,29 +193,48 @@ impl RangeScanCacheKey { } /// Cached result for one range scan. +#[derive(Debug)] +pub(crate) struct CachedBatchSlice { + batch: RecordBatch, + slice_lengths: Vec, +} + +impl CachedBatchSlice { + fn metadata_size(&self) -> usize { + self.slice_lengths.capacity() * mem::size_of::() + } +} + pub(crate) struct RangeScanCacheValue { - pub(crate) batches: Vec, - /// Precomputed size of all batches, accounting for shared dictionary values. + cached_batches: Vec, + /// Precomputed size of all compacted batches. estimated_batches_size: usize, } impl RangeScanCacheValue { - pub(crate) fn new(batches: Vec, estimated_batches_size: usize) -> Self { + pub(crate) fn new( + cached_batches: Vec, + estimated_batches_size: usize, + ) -> Self { Self { - batches, + cached_batches, estimated_batches_size, } } pub(crate) fn estimated_size(&self) -> usize { mem::size_of::() - + self.batches.capacity() * mem::size_of::() + + self.cached_batches.capacity() * mem::size_of::() + + self + .cached_batches + .iter() + .map(CachedBatchSlice::metadata_size) + .sum::() + self.estimated_batches_size } } /// Row groups and whether all sources are file-only for a partition range. -#[allow(dead_code)] pub(crate) struct PartitionRangeRowGroups { /// Sorted (file_id, row_group_index) pairs. pub(crate) row_groups: Vec<(FileId, i64)>, @@ -217,7 +242,6 @@ pub(crate) struct PartitionRangeRowGroups { } /// Collects (file_id, row_group_index) pairs from a partition range's row group indices. -#[allow(dead_code)] pub(crate) fn collect_partition_range_row_groups( stream_ctx: &StreamContext, part_range: &PartitionRange, @@ -244,11 +268,14 @@ pub(crate) fn collect_partition_range_row_groups( } /// Builds a cache key for the given partition range if it is eligible for caching. -#[allow(dead_code)] pub(crate) fn build_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, ) -> Option { + if !stream_ctx.input.cache_strategy.has_range_result_cache() { + return None; + } + let fingerprint = stream_ctx.scan_fingerprint.as_ref()?; // Dyn filters can change at runtime, so we can't cache when they're present. @@ -283,7 +310,6 @@ pub(crate) fn build_range_cache_key( }) } -#[allow(dead_code)] fn query_time_range_covers_partition_range( query_time_range: Option<&TimestampRange>, partition_time_range: FileTimeRange, @@ -297,117 +323,232 @@ fn query_time_range_covers_partition_range( } /// Returns a stream that replays cached record batches. -#[allow(dead_code)] pub(crate) fn cached_flat_range_stream(value: Arc) -> BoxedRecordBatchStream { - Box::pin(futures::stream::iter( - value.batches.clone().into_iter().map(Ok), - )) + Box::pin(try_stream! { + for cached_batch in &value.cached_batches { + let mut offset = 0; + for &len in &cached_batch.slice_lengths { + yield cached_batch.batch.slice(offset, len); + offset += len; + } + } + }) } -/// Returns true if two primary key dictionary arrays share the same underlying -/// values buffers by pointer comparison. -/// -/// The primary key column is always `DictionaryArray` with `Binary` values. -fn pk_values_ptr_eq(a: &DictionaryArray, b: &DictionaryArray) -> bool { - let a = a.values().as_binary::(); - let b = b.values().as_binary::(); - let values_eq = a.values().ptr_eq(b.values()) && a.offsets().ptr_eq(b.offsets()); - match (a.nulls(), b.nulls()) { - (Some(a), Some(b)) => values_eq && a.inner().ptr_eq(b.inner()), - (None, None) => values_eq, - _ => false, +enum CacheConcatCommand { + Compact(Vec), + Finish { + pending: Vec, + key: RangeScanCacheKey, + cache_strategy: CacheStrategy, + part_metrics: PartitionMetrics, + result_tx: Option>>>, + }, +} + +#[derive(Default)] +struct CacheConcatState { + cached_batches: Vec, + estimated_size: usize, +} + +impl CacheConcatState { + async fn compact( + &mut self, + batches: Vec, + limiter: &crate::cache::RangeResultMemoryLimiter, + ) -> Result<()> { + if batches.is_empty() { + return Ok(()); + } + + let input_size = batches + .iter() + .map(RecordBatch::get_array_memory_size) + .sum::(); + let _permit = limiter.acquire(input_size).await.map_err(|_| { + UnexpectedSnafu { + reason: "range result memory limiter is unexpectedly closed", + } + .build() + })?; + + let compacted = compact_record_batches(batches)?; + self.estimated_size += compacted.batch.get_array_memory_size(); + self.cached_batches.push(compacted); + Ok(()) + } + + fn finish(self) -> RangeScanCacheValue { + RangeScanCacheValue::new(self.cached_batches, self.estimated_size) } } -/// Buffers record batches for caching, tracking memory size while deduplicating -/// shared dictionary values across batches. -/// -/// Uses the primary key column as a proxy to detect dictionary sharing: if the PK -/// column's dictionary values are pointer-equal across batches, we assume all -/// dictionary columns share their values and deduct the total dictionary values size. -struct CacheBatchBuffer { +fn compact_record_batches(batches: Vec) -> Result { + debug_assert!(!batches.is_empty()); + + let slice_lengths = batches.iter().map(RecordBatch::num_rows).collect(); + build_cached_batch_slice(batches, slice_lengths) +} + +fn build_cached_batch_slice( batches: Vec, - /// Running total of batch memory. - total_size: usize, - /// The first batch's PK dictionary array, for pointer comparison. - /// `None` if no dictionary PK column exists or no batch has been added yet. - first_pk_dict: Option>, - /// Sum of `get_array_memory_size()` of all dictionary value arrays from the first batch. - total_dict_values_size: usize, - /// Whether the PK dictionary is still shared across all batches seen so far. - shared: bool, + slice_lengths: Vec, +) -> Result { + let batch = if batches.len() == 1 { + batches.into_iter().next().unwrap() + } else { + let schema = batches[0].schema(); + concat_batches(&schema, &batches).context(ComputeArrowSnafu)? + }; + + Ok(CachedBatchSlice { + batch, + slice_lengths, + }) +} + +async fn run_cache_concat_task( + mut rx: mpsc::UnboundedReceiver, + limiter: Arc, +) { + let mut state = CacheConcatState::default(); + + while let Some(cmd) = rx.recv().await { + match cmd { + CacheConcatCommand::Compact(batches) => { + if let Err(err) = state.compact(batches, &limiter).await { + warn!(err; "Failed to compact range cache batches"); + return; + } + } + CacheConcatCommand::Finish { + pending, + key, + cache_strategy, + part_metrics, + result_tx, + } => { + let result = state + .compact(pending, &limiter) + .await + .map(|()| state.finish()); + if let Err(err) = &result { + warn!(err; "Failed to finalize range cache batches"); + } + + let value = result.ok().map(Arc::new); + if let Some(value) = &value { + part_metrics + .inc_range_cache_size(key.estimated_size() + value.estimated_size()); + cache_strategy.put_range_result(key, value.clone()); + } + if let Some(tx) = result_tx { + let _ = tx.send(value); + } + return; + } + } + } +} + +struct CacheBatchBuffer { + buffered_batches: Vec, + buffered_rows: usize, + buffered_size: usize, + total_weight: usize, + sender: Option>, } impl CacheBatchBuffer { - fn new() -> Self { + fn new(cache_strategy: &CacheStrategy) -> Self { + let sender = cache_strategy.range_result_memory_limiter().map(|limiter| { + let (tx, rx) = mpsc::unbounded_channel(); + common_runtime::spawn_global(run_cache_concat_task(rx, limiter.clone())); + tx + }); + Self { - batches: Vec::new(), - total_size: 0, - first_pk_dict: None, - total_dict_values_size: 0, - shared: true, + buffered_batches: Vec::new(), + buffered_rows: 0, + buffered_size: 0, + total_weight: 0, + sender, } } - fn push(&mut self, batch: RecordBatch) { - if self.batches.is_empty() { - self.init_first_batch(&batch); - } else { - self.add_subsequent_batch(&batch); + fn push(&mut self, batch: RecordBatch) -> Result<()> { + if self.sender.is_none() { + return Ok(()); } - self.batches.push(batch); - } - fn init_first_batch(&mut self, batch: &RecordBatch) { - self.total_size += batch.get_array_memory_size(); - - let pk_col_idx = primary_key_column_index(batch.num_columns()); - let mut total_dict_values_size = 0; - for col_idx in 0..batch.num_columns() { - let col = batch.column(col_idx); - if let Some(dict) = col.as_any().downcast_ref::>() { - total_dict_values_size += dict.values().get_array_memory_size(); - if col_idx == pk_col_idx { - self.first_pk_dict = Some(dict.clone()); - } - } - } - self.total_dict_values_size = total_dict_values_size; - } - - fn add_subsequent_batch(&mut self, batch: &RecordBatch) { let batch_size = batch.get_array_memory_size(); - - if self.shared - && let Some(first_pk_dict) = &self.first_pk_dict - { - let pk_col_idx = primary_key_column_index(batch.num_columns()); - let col = batch.column(pk_col_idx); - if let Some(dict) = col.as_any().downcast_ref::>() - && pk_values_ptr_eq(first_pk_dict, dict) - { - // PK dict is shared, deduct all dict values sizes. - self.total_size += batch_size - self.total_dict_values_size; - return; - } - // Dictionary diverged. - self.shared = false; + self.total_weight += batch_size; + if self.total_weight > RANGE_CACHE_SKIP_BYTES { + self.buffered_batches.clear(); + self.buffered_rows = 0; + self.buffered_size = 0; + self.sender = None; + return Ok(()); } - self.total_size += batch_size; + self.buffered_rows += batch.num_rows(); + self.buffered_size += batch_size; + self.buffered_batches.push(batch); + + if self.buffered_rows > DEFAULT_READ_BATCH_SIZE + || self.buffered_size > RANGE_CACHE_COMPACT_THRESHOLD_BYTES + { + self.notify_compact(); + } + + Ok(()) } - fn estimated_batches_size(&self) -> usize { - self.total_size + fn notify_compact(&mut self) { + if self.buffered_batches.is_empty() || self.sender.is_none() { + return; + } + + let batches = mem::take(&mut self.buffered_batches); + self.buffered_rows = 0; + self.buffered_size = 0; + + let Some(sender) = &self.sender else { + return; + }; + if sender.send(CacheConcatCommand::Compact(batches)).is_err() { + self.sender = None; + } } - fn into_batches(self) -> Vec { - self.batches + fn finish( + mut self, + key: RangeScanCacheKey, + cache_strategy: CacheStrategy, + part_metrics: PartitionMetrics, + result_tx: Option>>>, + ) { + let Some(sender) = self.sender.take() else { + return; + }; + + if sender + .send(CacheConcatCommand::Finish { + pending: mem::take(&mut self.buffered_batches), + key, + cache_strategy, + part_metrics, + result_tx, + }) + .is_err() + { + self.sender = None; + } } } /// Wraps a stream to cache its output for future range cache hits. -#[allow(dead_code)] pub(crate) fn cache_flat_range_stream( mut stream: BoxedRecordBatchStream, cache_strategy: CacheStrategy, @@ -415,17 +556,13 @@ pub(crate) fn cache_flat_range_stream( part_metrics: PartitionMetrics, ) -> BoxedRecordBatchStream { Box::pin(try_stream! { - let mut buffer = CacheBatchBuffer::new(); + let mut buffer = CacheBatchBuffer::new(&cache_strategy); while let Some(batch) = stream.try_next().await? { - buffer.push(batch.clone()); + buffer.push(batch.clone())?; yield batch; } - let estimated_size = buffer.estimated_batches_size(); - let batches = buffer.into_batches(); - let value = Arc::new(RangeScanCacheValue::new(batches, estimated_size)); - part_metrics.inc_range_cache_size(key.estimated_size() + value.estimated_size()); - cache_strategy.put_range_result(key, value); + buffer.finish(key, cache_strategy, part_metrics, None); }) } @@ -486,10 +623,11 @@ mod tests { use common_time::Timestamp; use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_common::ScalarValue; use datafusion_expr::{Expr, col, lit}; use smallvec::smallvec; - use store_api::storage::FileId; + use store_api::storage::{FileId, RegionId}; use super::*; use crate::cache::CacheManager; @@ -508,6 +646,44 @@ mod tests { )) } + fn test_cache_context(strategy: &CacheStrategy) -> (RangeScanCacheKey, PartitionMetrics) { + let region_id = RegionId::new(1, 1); + let key = RangeScanCacheKey { + region_id, + row_groups: vec![], + scan: ScanRequestFingerprintBuilder { + read_column_ids: vec![], + read_column_types: vec![], + filters: vec![], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: false, + merge_mode: MergeMode::LastRow, + partition_expr_version: 0, + } + .build(), + }; + + let metrics_set = ExecutionPlanMetricsSet::new(); + let part_metrics = + PartitionMetrics::new(region_id, 0, "test", Instant::now(), false, &metrics_set); + + assert!(strategy.get_range_result(&key).is_none()); + (key, part_metrics) + } + + async fn finish_cache_batch_buffer( + buffer: CacheBatchBuffer, + key: RangeScanCacheKey, + cache_strategy: CacheStrategy, + part_metrics: PartitionMetrics, + ) -> Option> { + let (tx, rx) = oneshot::channel(); + buffer.finish(key, cache_strategy, part_metrics, Some(tx)); + rx.await.context(crate::error::RecvSnafu).ok().flatten() + } + async fn new_stream_context( filters: Vec, query_time_range: Option, @@ -687,169 +863,175 @@ mod tests { ); } - /// Creates a test schema with 5 columns where the primary key dictionary column - /// is at index 2 (`num_columns - 3`), matching the flat format layout. - /// - /// Layout: `[field0: Int64, field1: Int64, pk: Dictionary, ts: Int64, seq: Int64]` - fn dict_test_schema() -> Arc { + fn test_schema() -> Arc { use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; - Arc::new(Schema::new(vec![ - Field::new("field0", ArrowDataType::Int64, false), - Field::new("field1", ArrowDataType::Int64, false), - Field::new( - "pk", - ArrowDataType::Dictionary( - Box::new(ArrowDataType::UInt32), - Box::new(ArrowDataType::Binary), - ), - false, - ), - Field::new("ts", ArrowDataType::Int64, false), - Field::new("seq", ArrowDataType::Int64, false), - ])) + + Arc::new(Schema::new(vec![Field::new( + "value", + ArrowDataType::Int64, + false, + )])) } - /// Helper to create a record batch with a dictionary column at the primary key position. - fn make_dict_batch( - schema: Arc, - dict_values: &datatypes::arrow::array::BinaryArray, - keys: &[u32], - int_values: &[i64], - ) -> RecordBatch { - use datatypes::arrow::array::{Int64Array, UInt32Array}; + fn make_batch(values: &[i64]) -> RecordBatch { + use datatypes::arrow::array::Int64Array; - let key_array = UInt32Array::from(keys.to_vec()); - let dict_array: DictionaryArray = - DictionaryArray::new(key_array, Arc::new(dict_values.clone())); - let int_array = Int64Array::from(int_values.to_vec()); - let zeros = Int64Array::from(vec![0i64; int_values.len()]); RecordBatch::try_new( - schema, - vec![ - Arc::new(zeros.clone()), - Arc::new(int_array), - Arc::new(dict_array), - Arc::new(zeros.clone()), - Arc::new(zeros), - ], + test_schema(), + vec![Arc::new(Int64Array::from(values.to_vec()))], ) .unwrap() } - /// Computes the total `get_array_memory_size()` of all dictionary value arrays in a batch. - fn compute_total_dict_values_size(batch: &RecordBatch) -> usize { - batch - .columns() - .iter() - .filter_map(|col| { - col.as_any() - .downcast_ref::>() - .map(|dict| dict.values().get_array_memory_size()) - }) - .sum() - } - - #[test] - fn cache_batch_buffer_empty() { - let buffer = CacheBatchBuffer::new(); - assert_eq!(buffer.estimated_batches_size(), 0); - assert!(buffer.into_batches().is_empty()); - } - - #[test] - fn cache_batch_buffer_single_batch() { + fn make_large_binary_batch(rows: usize, bytes_per_row: usize) -> RecordBatch { use datatypes::arrow::array::BinaryArray; + use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; - let schema = dict_test_schema(); - let dict_values = BinaryArray::from_vec(vec![b"a", b"b", b"c"]); - let batch = make_dict_batch(schema, &dict_values, &[0, 1, 2], &[10, 20, 30]); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + ArrowDataType::Binary, + false, + )])); + let payload = vec![b'x'; bytes_per_row]; + let values = (0..rows).map(|_| payload.as_slice()).collect::>(); - let full_size = batch.get_array_memory_size(); - - let mut buffer = CacheBatchBuffer::new(); - buffer.push(batch); - assert_eq!(buffer.estimated_batches_size(), full_size); - assert_eq!(buffer.into_batches().len(), 1); + RecordBatch::try_new(schema, vec![Arc::new(BinaryArray::from_vec(values))]).unwrap() } #[test] - fn cache_batch_buffer_shared_dictionary() { - use datatypes::arrow::array::BinaryArray; + fn compact_record_batches_keeps_original_boundaries() { + let batches = vec![make_batch(&[1, 2]), make_batch(&[3]), make_batch(&[4, 5])]; - let schema = dict_test_schema(); - let dict_values = BinaryArray::from_vec(vec![b"alpha", b"beta", b"gamma"]); + let compacted = compact_record_batches(batches).unwrap(); - // Two batches sharing the same dictionary values array. - let batch1 = make_dict_batch(schema.clone(), &dict_values, &[0, 1], &[10, 20]); - let batch2 = make_dict_batch(schema, &dict_values, &[1, 2], &[30, 40]); + assert_eq!(compacted.batch.num_rows(), 5); + assert_eq!(compacted.slice_lengths, vec![2, 1, 2]); + } - let batch1_full = batch1.get_array_memory_size(); - let batch2_full = batch2.get_array_memory_size(); + #[tokio::test] + async fn cached_flat_range_stream_replays_original_batches() { + let value = Arc::new(RangeScanCacheValue::new( + vec![CachedBatchSlice { + batch: make_batch(&[1, 2, 3]), + slice_lengths: vec![2, 1], + }], + make_batch(&[1, 2, 3]).get_array_memory_size(), + )); - // The total dictionary values size that should be deduplicated for the second batch. - let dict_values_size = compute_total_dict_values_size(&batch2); + let replayed = cached_flat_range_stream(value) + .try_collect::>() + .await + .unwrap(); - let mut buffer = CacheBatchBuffer::new(); - buffer.push(batch1); - buffer.push(batch2); + assert_eq!(replayed.len(), 2); + assert_eq!(replayed[0].num_rows(), 2); + assert_eq!(replayed[1].num_rows(), 1); + } - // Second batch's dict values should not be counted again. + #[tokio::test] + async fn cache_batch_buffer_finishes_pending_batches() { + let strategy = test_cache_strategy(); + let batch = make_batch(&[1, 2, 3]); + let expected_size = batch.get_array_memory_size(); + let (key, part_metrics) = test_cache_context(&strategy); + + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch).unwrap(); + + let value = finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics) + .await + .unwrap(); + assert_eq!(value.cached_batches.len(), 1); + assert_eq!(value.cached_batches[0].slice_lengths, vec![3]); + assert_eq!(value.estimated_batches_size, expected_size); + assert!(Arc::ptr_eq( + &value, + &strategy.get_range_result(&key).unwrap() + )); + } + + #[tokio::test] + async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() { + let strategy = test_cache_strategy(); + let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); + let (key, part_metrics) = test_cache_context(&strategy); + + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch.clone()).unwrap(); + buffer.push(batch).unwrap(); + + assert_eq!(buffer.buffered_rows, 0); + assert!(buffer.buffered_batches.is_empty()); + + let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) + .await + .unwrap(); + assert_eq!(value.cached_batches.len(), 1); assert_eq!( - buffer.estimated_batches_size(), - batch1_full + batch2_full - dict_values_size + value.cached_batches[0].slice_lengths, + vec![ + DEFAULT_READ_BATCH_SIZE / 2 + 1, + DEFAULT_READ_BATCH_SIZE / 2 + 1 + ] ); - assert_eq!(buffer.into_batches().len(), 2); } - #[test] - fn cache_batch_buffer_non_shared_dictionary() { - use datatypes::arrow::array::BinaryArray; + #[tokio::test] + async fn cache_batch_buffer_compacts_when_buffered_size_exceeds_threshold() { + let strategy = test_cache_strategy(); + let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE, 4096); + let (key, part_metrics) = test_cache_context(&strategy); - let schema = dict_test_schema(); - let dict_values1 = BinaryArray::from_vec(vec![b"a", b"b"]); - let dict_values2 = BinaryArray::from_vec(vec![b"x", b"y"]); + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(large_batch.clone()).unwrap(); - let batch1 = make_dict_batch(schema.clone(), &dict_values1, &[0, 1], &[10, 20]); - let batch2 = make_dict_batch(schema, &dict_values2, &[0, 1], &[30, 40]); + assert_eq!(buffer.buffered_rows, 0); + assert!(buffer.buffered_batches.is_empty()); - let batch1_full = batch1.get_array_memory_size(); - let batch2_full = batch2.get_array_memory_size(); - - let mut buffer = CacheBatchBuffer::new(); - buffer.push(batch1); - buffer.push(batch2); - - // Different dictionaries: full size for both. - assert_eq!(buffer.estimated_batches_size(), batch1_full + batch2_full); - } - - #[test] - fn cache_batch_buffer_shared_then_diverged() { - use datatypes::arrow::array::BinaryArray; - - let schema = dict_test_schema(); - let shared_values = BinaryArray::from_vec(vec![b"a", b"b", b"c"]); - let different_values = BinaryArray::from_vec(vec![b"x", b"y"]); - - let batch1 = make_dict_batch(schema.clone(), &shared_values, &[0], &[1]); - let batch2 = make_dict_batch(schema.clone(), &shared_values, &[1], &[2]); - let batch3 = make_dict_batch(schema, &different_values, &[0], &[3]); - - let size1 = batch1.get_array_memory_size(); - let size2 = batch2.get_array_memory_size(); - let size3 = batch3.get_array_memory_size(); - - let dict_values_size = compute_total_dict_values_size(&batch2); - - let mut buffer = CacheBatchBuffer::new(); - buffer.push(batch1); - buffer.push(batch2); - buffer.push(batch3); - - // batch2 shares dict with batch1 (dedup), batch3 does not (full size). + let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) + .await + .unwrap(); + assert_eq!(value.cached_batches.len(), 1); assert_eq!( - buffer.estimated_batches_size(), - size1 + (size2 - dict_values_size) + size3 + value.cached_batches[0].slice_lengths, + vec![large_batch.num_rows()] + ); + } + + #[tokio::test] + async fn cache_batch_buffer_uses_compacted_size_for_weight() { + let strategy = test_cache_strategy(); + let batch1 = make_batch(&[1, 2]); + let batch2 = make_batch(&[3, 4]); + let (key, part_metrics) = test_cache_context(&strategy); + let expected = concat_batches(&test_schema(), &[batch1.clone(), batch2.clone()]) + .unwrap() + .get_array_memory_size(); + + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch1).unwrap(); + buffer.push(batch2).unwrap(); + + let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) + .await + .unwrap(); + assert_eq!(value.estimated_batches_size, expected); + } + + #[tokio::test] + async fn cache_batch_buffer_skips_cache_when_weight_exceeds_limit() { + let strategy = test_cache_strategy(); + let (key, part_metrics) = test_cache_context(&strategy); + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.total_weight = RANGE_CACHE_SKIP_BYTES; + + buffer.push(make_batch(&[1])).unwrap(); + + assert!(buffer.sender.is_none()); + assert!( + finish_cache_batch_buffer(buffer, key, strategy, part_metrics) + .await + .is_none() ); } } diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 8fc946b3d3..80563f32a9 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1258,13 +1258,25 @@ pub(crate) fn should_split_flat_batches_for_merge( // This is a file range. let file_index = index.index - stream_ctx.input.num_memtables(); let file = &stream_ctx.input.files[file_index]; - if file.meta_ref().num_rows < SPLIT_ROW_THRESHOLD || file.meta_ref().num_series == 0 { + let file_meta = file.meta_ref(); + if file_meta.level == 0 { + // Always split level 0 files. + num_files_to_split += 1; + continue; + } else if file_meta.num_rows < SPLIT_ROW_THRESHOLD || file_meta.num_series == 0 { // If the file doesn't have enough rows, or the number of series is unavailable, skips it. continue; } - debug_assert!(file.meta_ref().num_rows > 0); - if !can_split_series(file.meta_ref().num_rows, file.meta_ref().num_series) { + debug_assert!(file_meta.num_rows > 0); + if !can_split_series(file_meta.num_rows, file_meta.num_series) { // We can't split batches in a file. + common_telemetry::trace!( + "Can't split series for file {}, level: {}, num_rows: {}, num_series: {}", + file_meta.file_id, + file_meta.level, + file_meta.num_rows, + file_meta.num_series, + ); return None; } else { num_files_to_split += 1; @@ -1310,14 +1322,108 @@ pub(crate) fn compute_parallel_channel_size(estimated_rows_per_batch: usize) -> size.clamp(2, 64) } +/// Computes the average estimated rows per batch across multiple range readers. +pub(crate) fn compute_average_batch_size( + estimated_rows_per_batch: impl IntoIterator, +) -> usize { + let mut total = 0usize; + let mut count = 0usize; + for size in estimated_rows_per_batch { + total += size; + count += 1; + } + + if count == 0 { + return DEFAULT_READ_BATCH_SIZE; + } + + (total / count).clamp(1, DEFAULT_READ_BATCH_SIZE) +} + fn can_split_series(num_rows: u64, num_series: u64) -> bool { - assert!(num_series > 0); - assert!(num_rows > 0); + if num_rows == 0 || num_series == 0 { + return false; + } // It doesn't have too many series or it will have enough rows for each batch. num_series < NUM_SERIES_THRESHOLD || num_rows / num_series >= BATCH_SIZE_THRESHOLD } +#[cfg(test)] +mod split_tests { + use std::sync::Arc; + + use common_time::Timestamp; + use smallvec::smallvec; + use store_api::storage::FileId; + + use super::*; + use crate::read::projection::ProjectionMapper; + use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex}; + use crate::read::scan_region::{ScanInput, StreamContext}; + use crate::sst::file::FileHandle; + use crate::test_util::memtable_util::metadata_with_primary_key; + use crate::test_util::scheduler_util::SchedulerEnv; + use crate::test_util::sst_util::sst_file_handle_with_file_id; + + async fn new_stream_context_with_files(files: Vec) -> StreamContext { + let env = SchedulerEnv::new().await; + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let mapper = ProjectionMapper::new(&metadata, [0, 2, 3].into_iter()).unwrap(); + let input = ScanInput::new(env.access_layer.clone(), mapper).with_files(files); + + StreamContext { + input, + ranges: vec![], + scan_fingerprint: None, + query_start: std::time::Instant::now(), + } + } + + fn single_file_range_meta() -> RangeMeta { + RangeMeta { + time_range: ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(1000), + ), + indices: smallvec![SourceIndex { + index: 0, + num_row_groups: 1, + }], + row_group_indices: smallvec![RowGroupIndex { + index: 0, + row_group_index: 0, + }], + num_rows: 1024, + } + } + + #[tokio::test] + async fn should_split_level_zero_file_even_when_series_stats_are_missing() { + let mut file = sst_file_handle_with_file_id(FileId::random(), 0, 1000) + .meta_ref() + .clone(); + file.level = 0; + file.num_rows = DEFAULT_ROW_GROUP_SIZE as u64; + file.num_row_groups = 1; + file.num_series = 0; + + let file = FileHandle::new(file, crate::test_util::new_noop_file_purger()); + let stream_ctx = Arc::new(new_stream_context_with_files(vec![file]).await); + + assert!( + should_split_flat_batches_for_merge(&stream_ctx, &single_file_range_meta()).is_some() + ); + } + + #[test] + fn can_split_series_returns_false_for_zero_inputs() { + assert!(!can_split_series(0, 1)); + assert!(!can_split_series(1, 0)); + assert!(!can_split_series(0, 0)); + } +} + /// Creates a new [ReaderFilterMetrics] with optional apply metrics initialized /// based on the `explain_verbose` flag. fn new_filter_metrics(explain_verbose: bool) -> ReaderFilterMetrics { @@ -1653,6 +1759,7 @@ mod tests { let meta = FileMeta { region_id: RegionId::new(123, 456), file_id: Default::default(), + level: 1, time_range: ( Timestamp::new_millisecond(0), Timestamp::new_millisecond(1000), @@ -1816,4 +1923,26 @@ mod tests { compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE * 2) ); } + + #[test] + fn test_compute_average_batch_size_uses_arithmetic_mean() { + assert_eq!(24, compute_average_batch_size([16, 24, 32])); + } + + #[test] + fn test_compute_average_batch_size_clamps_values() { + assert_eq!( + DEFAULT_READ_BATCH_SIZE, + compute_average_batch_size([DEFAULT_READ_BATCH_SIZE, DEFAULT_READ_BATCH_SIZE * 2]) + ); + assert_eq!(1, compute_average_batch_size([0, 1])); + } + + #[test] + fn test_compute_average_batch_size_falls_back_when_empty() { + assert_eq!( + DEFAULT_READ_BATCH_SIZE, + compute_average_batch_size(std::iter::empty()) + ); + } } diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 15ab435425..932d382834 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -41,6 +41,9 @@ use crate::read::flat_merge::FlatMergeReader; use crate::read::last_row::FlatLastRowReader; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RangeMeta; +use crate::read::range_cache::{ + build_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, +}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, SplitRecordBatchStream, compute_parallel_channel_size, @@ -181,19 +184,22 @@ impl SeqScan { sources, None, None, + false, compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE), ) .await } - /// Builds a flat reader to read sources that returns RecordBatch. If `semaphore` is provided, reads sources in parallel - /// if possible. + /// Builds a flat reader to read sources that returns RecordBatch. + /// If `semaphore` is provided, reads sources in parallel if possible. + /// If `skip_dedup` is true, the merged stream is returned without applying flat dedup. #[tracing::instrument(level = tracing::Level::DEBUG, skip_all)] pub(crate) async fn build_flat_reader_from_sources( stream_ctx: &StreamContext, mut sources: Vec, semaphore: Option>, part_metrics: Option<&PartitionMetrics>, + skip_dedup: bool, channel_size: usize, ) -> Result { if let Some(semaphore) = semaphore.as_ref() { @@ -215,7 +221,7 @@ impl SeqScan { FlatMergeReader::new(schema, sources, DEFAULT_READ_BATCH_SIZE, metrics_reporter) .await?; - let dedup = !stream_ctx.input.append_mode; + let dedup = !skip_dedup && !stream_ctx.input.append_mode; let dedup_metrics_reporter = part_metrics.map(|m| m.dedup_metrics_reporter()); let reader = if dedup { match stream_ctx.input.merge_mode { @@ -253,6 +259,62 @@ impl SeqScan { Ok(reader) } + /// Builds a flat read stream for one partition range. + pub(crate) async fn build_flat_partition_range_read( + stream_ctx: &Arc, + part_range: &PartitionRange, + compaction: bool, + part_metrics: &PartitionMetrics, + partition_pruner: Arc, + file_scan_semaphore: Option>, + merge_semaphore: Option>, + ) -> Result<(BoxedRecordBatchStream, usize)> { + let cache_key = build_range_cache_key(stream_ctx, part_range); + + if let Some(key) = cache_key.as_ref() { + if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { + part_metrics.inc_range_cache_hit(); + return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); + } + part_metrics.inc_range_cache_miss(); + } + + let mut sources = Vec::new(); + let split_batch_size = build_flat_sources( + stream_ctx, + part_range, + compaction, + part_metrics, + partition_pruner, + &mut sources, + file_scan_semaphore, + ) + .await?; + let estimated_rows_per_batch = split_batch_size.unwrap_or(DEFAULT_READ_BATCH_SIZE); + let channel_size = compute_parallel_channel_size(estimated_rows_per_batch); + let stream = Self::build_flat_reader_from_sources( + stream_ctx, + sources, + merge_semaphore, + Some(part_metrics), + false, + channel_size, + ) + .await?; + + let stream = match cache_key { + Some(key) => cache_flat_range_stream( + stream, + stream_ctx.input.cache_strategy.clone(), + key, + part_metrics.clone(), + ), + None => stream, + }; + + Ok((stream, estimated_rows_per_batch)) + } + /// Scans the given partition when the part list is set properly. /// Otherwise the returned stream might not contains any data. fn scan_partition_impl( @@ -331,23 +393,16 @@ impl SeqScan { // Scans each part. for part_range in partition_ranges { - let mut sources = Vec::new(); - let split_batch_size = build_flat_sources( + let (mut reader, _) = Self::build_flat_partition_range_read( &stream_ctx, &part_range, compaction, &part_metrics, partition_pruner.clone(), - &mut sources, file_scan_semaphore.clone(), - ).await?; - - let channel_size = compute_parallel_channel_size( - split_batch_size.unwrap_or(DEFAULT_READ_BATCH_SIZE), - ); - let mut reader = - Self::build_flat_reader_from_sources(&stream_ctx, sources, semaphore.clone(), Some(&part_metrics), channel_size) - .await?; + semaphore.clone(), + ) + .await?; let mut metrics = ScannerMetrics { scan_cost: fetch_start.elapsed(), diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index bf7ed072ab..7883c1d553 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -41,18 +41,18 @@ use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError}; use tokio::sync::mpsc::{self, Receiver, Sender}; use crate::error::{ - Error, InvalidSenderSnafu, PartitionOutOfRangeSnafu, Result, ScanMultiTimesSnafu, + Error, InvalidSenderSnafu, JoinSnafu, PartitionOutOfRangeSnafu, Result, ScanMultiTimesSnafu, ScanSeriesSnafu, TooManyFilesToReadSnafu, }; use crate::read::ScannerMetrics; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ - PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_parallel_channel_size, + PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_average_batch_size, + compute_parallel_channel_size, }; -use crate::read::seq_scan::{SeqScan, build_flat_sources}; +use crate::read::seq_scan::SeqScan; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; -use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::PrimaryKeyArray; @@ -227,7 +227,8 @@ impl SeriesScan { let (senders, receivers) = new_channel_list(self.properties.num_partitions()); let mut distributor = SeriesDistributor { stream_ctx: self.stream_ctx.clone(), - semaphore: Some(Arc::new(Semaphore::new(self.properties.num_partitions()))), + range_semaphore: Some(Arc::new(Semaphore::new(self.properties.num_partitions()))), + final_merge_semaphore: Some(Arc::new(Semaphore::new(self.properties.num_partitions()))), partitions: self.properties.partitions.clone(), pruner: self.pruner.clone(), senders, @@ -420,8 +421,13 @@ impl SeriesScan { struct SeriesDistributor { /// Context for the scan stream. stream_ctx: Arc, - /// Optional semaphore for limiting the number of concurrent scans. - semaphore: Option>, + /// Semaphore for file scanning and range-level merging. + range_semaphore: Option>, + /// Semaphore for the final merge across all range streams. + /// Must be separate from `range_semaphore` to avoid deadlock: final merge tasks + /// hold a permit while waiting for data from range-level merge tasks, which also + /// need permits to produce data. + final_merge_semaphore: Option>, /// Partition ranges to scan. partitions: Vec>, /// Shared pruner for file range building. @@ -483,36 +489,57 @@ impl SeriesDistributor { // build part cost. let mut fetch_start = Instant::now(); - // Scans all parts. - let mut sources = Vec::with_capacity(self.partitions.len()); - let mut min_batch_size: Option = None; + // Builds one deduped stream per partition range, then merges across ranges. + let build_start = Instant::now(); + let mut tasks = Vec::new(); for partition in &self.partitions { - sources.reserve(partition.len()); for part_range in partition { - let split_batch_size = build_flat_sources( - &self.stream_ctx, - part_range, - false, - &part_metrics, - partition_pruner.clone(), - &mut sources, - self.semaphore.clone(), - ) - .await?; - if let Some(size) = split_batch_size { - min_batch_size = Some(min_batch_size.map_or(size, |cur| cur.min(size))); - } + let stream_ctx = self.stream_ctx.clone(); + let part_range = *part_range; + let part_metrics = part_metrics.clone(); + let partition_pruner = partition_pruner.clone(); + let file_scan_semaphore = self.range_semaphore.clone(); + let merge_semaphore = self.range_semaphore.clone(); + tasks.push(common_runtime::spawn_global(async move { + SeqScan::build_flat_partition_range_read( + &stream_ctx, + &part_range, + false, + &part_metrics, + partition_pruner, + file_scan_semaphore, + merge_semaphore, + ) + .await + })); } } - - // Builds a flat reader that merge sources from all parts. + let mut range_streams = Vec::with_capacity(tasks.len()); + let mut estimated_batch_sizes = Vec::with_capacity(tasks.len()); + for task in tasks { + let (stream, estimated_batch_size) = task.await.context(JoinSnafu)??; + range_streams.push(stream); + estimated_batch_sizes.push(estimated_batch_size); + } let channel_size = - compute_parallel_channel_size(min_batch_size.unwrap_or(DEFAULT_READ_BATCH_SIZE)); + compute_parallel_channel_size(compute_average_batch_size(estimated_batch_sizes)); + common_telemetry::debug!( + "SeriesDistributor built {} range_streams, region: {}, build cost: {:?}, channel_size: {}", + range_streams.len(), + self.stream_ctx.input.region_metadata().region_id, + build_start.elapsed(), + channel_size, + ); + + // Each partition range stream is already deduped, so skip dedup here. + // Use a separate semaphore for the final merge to avoid deadlock with + // range-level merge tasks that share the range_semaphore. let mut reader = SeqScan::build_flat_reader_from_sources( &self.stream_ctx, - sources, - self.semaphore.clone(), + range_streams, + self.final_merge_semaphore.clone(), Some(&part_metrics), + true, channel_size, ) .await?; From 57f19212534b2329b1418fd4297b3acf58287eb4 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 13 Apr 2026 17:04:02 +0800 Subject: [PATCH 014/120] feat: propagate staging leader through lease and heartbeat (#7950) * feat(mito): expose staging leader role state * fix(region): clear staging metadata on leader exit * feat: propagate staging leader role through heartbeat and metasrv * chore: update comments Signed-off-by: WenyXu * fix(region): unify staging exit role transitions * chore: update proto Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/common/meta/src/datanode.rs | 23 ++ src/datanode/src/alive_keeper.rs | 138 +++++++++ src/datanode/src/heartbeat.rs | 6 +- src/datanode/src/region_server.rs | 1 + .../handler/collect_cluster_info_handler.rs | 2 +- .../handler/collect_leader_region_handler.rs | 2 +- .../src/handler/persist_stats_handler.rs | 5 +- .../src/handler/region_lease_handler.rs | 59 ++++ src/meta-srv/src/region/lease_keeper.rs | 124 +++++++- src/mito2/src/engine.rs | 10 +- .../src/engine/apply_staging_manifest_test.rs | 2 +- src/mito2/src/engine/set_role_state_test.rs | 266 +++++++++++++++++- src/mito2/src/engine/staging_test.rs | 2 +- src/mito2/src/region.rs | 173 +++++++++--- src/mito2/src/region/opener.rs | 5 +- src/mito2/src/worker/handle_apply_staging.rs | 2 +- src/mito2/src/worker/handle_enter_staging.rs | 10 +- src/store-api/src/region_engine.rs | 12 +- 20 files changed, 782 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 872095752b..4f6339d83a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5681,7 +5681,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=092ba1d01e2da676dca66cca7eebb55009da8ef8#092ba1d01e2da676dca66cca7eebb55009da8ef8" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=26a50f4069f50c37d65b45e0d39ae0cb42de5425#26a50f4069f50c37d65b45e0d39ae0cb42de5425" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -5691,7 +5691,6 @@ dependencies = [ "strum_macros 0.25.3", "tonic 0.14.2", "tonic-prost", - "tonic-prost-build", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 227608bf64..34e10d9173 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -154,7 +154,7 @@ etcd-client = { version = "0.17", features = [ fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "092ba1d01e2da676dca66cca7eebb55009da8ef8" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "26a50f4069f50c37d65b45e0d39ae0cb42de5425" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index 8b521d8e43..d6c6229801 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -573,4 +573,27 @@ mod tests { let region_num = stat_val.region_num().unwrap(); assert_eq!(2, region_num); } + + #[test] + fn test_region_stat_from_heartbeat_preserves_staging_leader_role() { + let request = HeartbeatRequest { + header: Some(RequestHeader::default()), + peer: Some(api::v1::meta::Peer { + id: 1, + addr: "127.0.0.1:3001".to_string(), + }), + region_stats: vec![api::v1::meta::RegionStat { + region_id: RegionId::new(1024, 1).as_u64(), + engine: "mito".to_string(), + role: api::v1::meta::RegionRole::StagingLeader.into(), + ..Default::default() + }], + ..Default::default() + }; + + let stat = Stat::try_from(&request).unwrap(); + + assert_eq!(stat.region_stats.len(), 1); + assert_eq!(stat.region_stats[0].role, RegionRole::StagingLeader); + } } diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index 57f4e00aa2..dbf99fdb28 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -503,6 +503,7 @@ mod test { use mito2::config::MitoConfig; use mito2::test_util::{CreateRequestBuilder, TestEnv}; use store_api::region_engine::RegionEngine; + use store_api::region_request::{EnterStagingRequest, StagingPartitionDirective}; use super::*; use crate::tests::mock_region_server; @@ -621,4 +622,141 @@ mod test { > Instant::now() + Duration::from_millis(heartbeat_interval_millis * 4) ); } + + #[tokio::test(flavor = "multi_thread")] + async fn renew_staging_leader_keeps_region_in_staging() { + let mut region_server = mock_region_server(); + let mut engine_env = TestEnv::with_prefix("region-alive-keeper-staging").await; + let engine = engine_env.create_engine(MitoConfig::default()).await; + let engine = Arc::new(engine); + region_server.register_engine(engine.clone()); + + let alive_keeper = Arc::new(RegionAliveKeeper::new( + region_server.clone(), + None, + Duration::from_millis(100), + )); + + let region_id = RegionId::new(1024, 2); + region_server + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + region_server + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + alive_keeper.register_region(region_id).await; + alive_keeper + .renew_region_leases( + &[GrantedRegion { + region_id: region_id.as_u64(), + role: api::v1::meta::RegionRole::StagingLeader.into(), + extensions: HashMap::new(), + }], + Instant::now() + Duration::from_millis(3000), + ) + .await; + + assert_eq!(engine.role(region_id).unwrap(), RegionRole::StagingLeader); + } + + #[tokio::test(flavor = "multi_thread")] + async fn renew_staging_leader_exit_into_leader() { + common_telemetry::init_default_ut_logging(); + let mut region_server = mock_region_server(); + let mut engine_env = TestEnv::with_prefix("region-alive-keeper-staging-exit").await; + let engine = engine_env.create_engine(MitoConfig::default()).await; + let engine = Arc::new(engine); + region_server.register_engine(engine.clone()); + + let alive_keeper = Arc::new(RegionAliveKeeper::new( + region_server.clone(), + None, + Duration::from_millis(100), + )); + + let region_id = RegionId::new(1024, 2); + region_server + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + region_server + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + alive_keeper.register_region(region_id).await; + alive_keeper + .renew_region_leases( + &[GrantedRegion { + region_id: region_id.as_u64(), + role: api::v1::meta::RegionRole::Leader.into(), + extensions: HashMap::new(), + }], + Instant::now() + Duration::from_millis(3000), + ) + .await; + + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(engine.role(region_id).unwrap(), RegionRole::Leader); + } + + #[tokio::test(flavor = "multi_thread")] + async fn renew_staging_leader_does_not_promote_normal_leader_into_staging() { + let mut region_server = mock_region_server(); + let mut engine_env = TestEnv::with_prefix("region-alive-keeper-non-staging").await; + let engine = engine_env.create_engine(MitoConfig::default()).await; + let engine = Arc::new(engine); + region_server.register_engine(engine.clone()); + + let alive_keeper = Arc::new(RegionAliveKeeper::new( + region_server.clone(), + None, + Duration::from_millis(100), + )); + + let region_id = RegionId::new(1024, 4); + region_server + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + region_server + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + alive_keeper.register_region(region_id).await; + alive_keeper + .renew_region_leases( + &[GrantedRegion { + region_id: region_id.as_u64(), + role: api::v1::meta::RegionRole::StagingLeader.into(), + extensions: HashMap::new(), + }], + Instant::now() + Duration::from_millis(3000), + ) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(engine.role(region_id).unwrap(), RegionRole::Leader); + } } diff --git a/src/datanode/src/heartbeat.rs b/src/datanode/src/heartbeat.rs index be662dfe94..fe8866b7f9 100644 --- a/src/datanode/src/heartbeat.rs +++ b/src/datanode/src/heartbeat.rs @@ -148,9 +148,9 @@ impl HeartbeatTask { let mut follower_region_lease_count = 0; for lease in &lease.regions { match lease.role() { - RegionRole::Leader | RegionRole::DowngradingLeader => { - leader_region_lease_count += 1 - } + RegionRole::Leader + | RegionRole::StagingLeader + | RegionRole::DowngradingLeader => leader_region_lease_count += 1, RegionRole::Follower => follower_region_lease_count += 1, } } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index ec10691bea..aa3ffbfe3a 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -360,6 +360,7 @@ impl RegionServer { engine.role(region_id).map(|role| match role { RegionRole::Follower => false, RegionRole::Leader => true, + RegionRole::StagingLeader => true, RegionRole::DowngradingLeader => true, }) }) diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index c96229f9cf..3fc785a1cb 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -129,7 +129,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { let leader_regions = stat .region_stats .iter() - .filter(|s| s.role == RegionRole::Leader) + .filter(|s| matches!(s.role, RegionRole::Leader | RegionRole::StagingLeader)) .count(); let follower_regions = stat.region_stats.len() - leader_regions; diff --git a/src/meta-srv/src/handler/collect_leader_region_handler.rs b/src/meta-srv/src/handler/collect_leader_region_handler.rs index ddb4cd0ea3..95b03e3341 100644 --- a/src/meta-srv/src/handler/collect_leader_region_handler.rs +++ b/src/meta-srv/src/handler/collect_leader_region_handler.rs @@ -40,7 +40,7 @@ impl HeartbeatHandler for CollectLeaderRegionHandler { let mut key_values = Vec::with_capacity(current_stat.region_stats.len()); for stat in current_stat.region_stats.iter() { - if stat.role != RegionRole::Leader { + if !matches!(stat.role, RegionRole::Leader | RegionRole::StagingLeader) { continue; } diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index 75281f982a..d863070225 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -121,7 +121,10 @@ fn to_persisted_if_leader( datanode_id: DatanodeId, timestamp_millis: i64, ) -> Option<(Row, PersistedRegionStat)> { - if matches!(region_stat.role, RegionRole::Leader) { + if matches!( + region_stat.role, + RegionRole::Leader | RegionRole::StagingLeader + ) { let persisted_region_stat = last_persisted_region_stats.get(®ion_stat.id).map(|s| *s); Some(( compute_persist_region_stat( diff --git a/src/meta-srv/src/handler/region_lease_handler.rs b/src/meta-srv/src/handler/region_lease_handler.rs index fef84ef0db..c6c1d44521 100644 --- a/src/meta-srv/src/handler/region_lease_handler.rs +++ b/src/meta-srv/src/handler/region_lease_handler.rs @@ -398,6 +398,65 @@ mod test { assert_eq!(acc.inactive_region_ids, HashSet::from([no_exist_region_id])); } + #[tokio::test] + async fn test_handle_staging_leader() { + let datanode_id = 1; + let region_number = 1u32; + let table_id = 10; + let region_id = RegionId::new(table_id, region_number); + let peer = Peer::empty(datanode_id); + let table_info = new_test_table_info(table_id); + + let region_routes = vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(peer.clone()), + leader_state: Some(LeaderState::Staging), + ..Default::default() + }]; + + let keeper = new_test_keeper(); + let table_metadata_manager = keeper.table_metadata_manager(); + + table_metadata_manager + .create_table_metadata( + table_info, + TableRouteValue::physical(region_routes), + HashMap::default(), + ) + .await + .unwrap(); + + let builder = MetasrvBuilder::new(); + let metasrv = builder.build().await.unwrap(); + let ctx = &mut metasrv.new_ctx(); + + let req = HeartbeatRequest { + duration_since_epoch: 1234, + ..Default::default() + }; + + let acc = &mut HeartbeatAccumulator::default(); + acc.stat = Some(Stat { + id: peer.id, + region_stats: vec![new_empty_region_stat(region_id, RegionRole::StagingLeader)], + ..Default::default() + }); + + let handler = RegionLeaseHandler::new( + default_distributed_time_constants().region_lease.as_secs(), + table_metadata_manager.clone(), + Default::default(), + None, + ); + + handler.handle(&req, ctx, acc).await.unwrap(); + + assert_region_lease( + acc, + vec![GrantedRegion::new(region_id, RegionRole::StagingLeader)], + ); + } + fn assert_region_lease(acc: &HeartbeatAccumulator, expected: Vec) { let region_lease = acc.region_lease.as_ref().unwrap().clone(); let granted: Vec = region_lease diff --git a/src/meta-srv/src/region/lease_keeper.rs b/src/meta-srv/src/region/lease_keeper.rs index 6d282fb49f..ac9f7d71b9 100644 --- a/src/meta-srv/src/region/lease_keeper.rs +++ b/src/meta-srv/src/region/lease_keeper.rs @@ -63,7 +63,9 @@ fn renew_region_lease_via_region_route( if let Some(leader) = ®ion_route.leader_peer && leader.id == datanode_id { - let region_role = if region_route.is_leader_downgrading() { + let region_role = if region_route.is_leader_staging() { + RegionRole::StagingLeader + } else if region_route.is_leader_downgrading() { RegionRole::DowngradingLeader } else { RegionRole::Leader @@ -313,6 +315,12 @@ mod tests { renew_region_lease_via_region_route(®ion_route, leader_peer_id, region_id), Some((region_id, RegionRole::DowngradingLeader)) ); + + region_route.leader_state = Some(LeaderState::Staging); + assert_eq!( + renew_region_lease_via_region_route(®ion_route, leader_peer_id, region_id), + Some((region_id, RegionRole::StagingLeader)) + ); } #[tokio::test] @@ -581,4 +589,118 @@ mod tests { ); } } + + #[tokio::test] + async fn test_renew_region_leases_reported_staging_expected_leader() { + let table_id = 1024; + let table_info: TableInfo = new_test_table_info(table_id); + + let region_id = RegionId::new(table_id, 1); + let leader_peer_id = 1024; + let region_route = RegionRouteBuilder::default() + .region(Region::new_test(region_id)) + .leader_peer(Peer::empty(leader_peer_id)) + .build() + .unwrap(); + + let keeper = new_test_keeper(); + let table_metadata_manager = keeper.table_metadata_manager(); + table_metadata_manager + .create_table_metadata( + table_info, + TableRouteValue::physical(vec![region_route]), + HashMap::default(), + ) + .await + .unwrap(); + + let RenewRegionLeasesResponse { + non_exists, + renewed, + } = keeper + .renew_region_leases(leader_peer_id, &[(region_id, RegionRole::StagingLeader)]) + .await + .unwrap(); + + assert!(non_exists.is_empty()); + assert_eq!( + renewed, + HashMap::from([( + region_id, + RegionLeaseInfo::from((region_id, RegionRole::Leader)) + )]) + ); + } + + #[tokio::test] + async fn test_renew_region_leases_reported_staging_expected_staging() { + let table_id = 1024; + let table_info: TableInfo = new_test_table_info(table_id); + + let region_id = RegionId::new(table_id, 1); + let leader_peer_id = 1024; + let region_route = RegionRouteBuilder::default() + .region(Region::new_test(region_id)) + .leader_peer(Peer::empty(leader_peer_id)) + .leader_state(LeaderState::Staging) + .build() + .unwrap(); + + let keeper = new_test_keeper(); + let table_metadata_manager = keeper.table_metadata_manager(); + table_metadata_manager + .create_table_metadata( + table_info, + TableRouteValue::physical(vec![region_route]), + HashMap::default(), + ) + .await + .unwrap(); + + let RenewRegionLeasesResponse { + non_exists, + renewed, + } = keeper + .renew_region_leases(leader_peer_id, &[(region_id, RegionRole::StagingLeader)]) + .await + .unwrap(); + + assert!(non_exists.is_empty()); + assert_eq!( + renewed, + HashMap::from([( + region_id, + RegionLeaseInfo::from((region_id, RegionRole::StagingLeader)) + )]) + ); + } + + #[tokio::test] + async fn test_renew_region_leases_operating_region_preserves_reported_role() { + let keeper = new_test_keeper(); + let datanode_id = 1024; + let region_id = RegionId::new(2048, 1); + + let _guard = keeper + .memory_region_keeper + .register(datanode_id, region_id) + .unwrap(); + + let RenewRegionLeasesResponse { + non_exists, + renewed, + } = keeper + .renew_region_leases(datanode_id, &[(region_id, RegionRole::StagingLeader)]) + .await + .unwrap(); + + assert!(non_exists.is_empty()); + assert_eq!( + renewed, + HashMap::from([( + region_id, + RegionLeaseInfo::operating(region_id, RegionRole::StagingLeader) + )]) + ); + } } diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index d006067f0d..5bd1002581 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -1114,13 +1114,9 @@ impl EngineInner { } fn role(&self, region_id: RegionId) -> Option { - self.workers.get_region(region_id).map(|region| { - if region.is_follower() { - RegionRole::Follower - } else { - RegionRole::Leader - } - }) + self.workers + .get_region(region_id) + .map(|region| region.region_role()) } } diff --git a/src/mito2/src/engine/apply_staging_manifest_test.rs b/src/mito2/src/engine/apply_staging_manifest_test.rs index a82fcfe049..efa0713cfc 100644 --- a/src/mito2/src/engine/apply_staging_manifest_test.rs +++ b/src/mito2/src/engine/apply_staging_manifest_test.rs @@ -333,7 +333,7 @@ async fn test_apply_staging_manifest_success_with_format(flat_format: bool) { let staging_manifest = region.manifest_ctx.staging_manifest().await; assert!(staging_manifest.is_none()); // The staging partition expr should be cleared. - assert!(region.staging_partition_info.lock().unwrap().is_none()); + assert!(region.manifest_ctx.staging_partition_info().is_none()); // The staging manifest directory should be empty. let data_home = env.data_home(); let region_dir = format!("{}/data/test/1_0000000001", data_home.display()); diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 4fb15ab7fe..40e03b063a 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -19,7 +19,9 @@ use store_api::region_engine::{ RegionEngine, RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; -use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::region_request::{ + EnterStagingRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective, +}; use store_api::storage::RegionId; use crate::config::MitoConfig; @@ -241,12 +243,14 @@ async fn test_unified_state_transitions_with_format(flat_format: bool) { .await .unwrap(); assert_success_response(&result, 0); + assert_eq!(engine.role(region_id), Some(RegionRole::StagingLeader)); let result = engine .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) .await .unwrap(); assert_success_response(&result, 0); + assert_eq!(engine.role(region_id), Some(RegionRole::Leader)); // Leader -> StagingLeader -> Follower (exit staging via demotion) engine @@ -259,6 +263,7 @@ async fn test_unified_state_transitions_with_format(flat_format: bool) { .await .unwrap(); assert_success_response(&result, 0); + assert_eq!(engine.role(region_id), Some(RegionRole::Follower)); // Note: Direct Follower -> Leader promotion is no longer allowed // Use existing set_region_role method for follower -> leader promotion @@ -277,6 +282,7 @@ async fn test_unified_state_transitions_with_format(flat_format: bool) { .await .unwrap(); assert_success_response(&result, 0); + assert_eq!(engine.role(region_id), Some(RegionRole::DowngradingLeader)); // Note: Direct DowngradingLeader -> Leader is no longer allowed // Use existing set_region_role method for downgrading -> leader promotion @@ -325,6 +331,264 @@ async fn test_restricted_state_transitions() { test_restricted_state_transitions_with_format(true).await; } +#[tokio::test] +async fn test_direct_set_region_role_staging_leader_is_noop() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .set_region_role(region_id, RegionRole::StagingLeader) + .unwrap(); + + assert_eq!(engine.role(region_id), Some(RegionRole::Leader)); + + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + engine + .set_region_role(region_id, RegionRole::StagingLeader) + .unwrap(); + + assert_eq!(engine.role(region_id), Some(RegionRole::Follower)); +} + +#[tokio::test] +async fn test_direct_set_region_role_exits_staging_state_only() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::StagingLeader)); + assert!( + engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .staging_partition_info() + .is_some() + ); + + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::Leader)); + assert!( + engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .staging_partition_info() + .is_none() + ); +} + +#[tokio::test] +async fn test_set_region_role_can_exit_staging_to_leader() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::StagingLeader) + .await + .unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::StagingLeader)); + + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + assert_eq!(engine.role(region_id), Some(RegionRole::Leader)); + assert!( + engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .staging_partition_info() + .is_none() + ); +} + +#[tokio::test] +async fn test_set_region_role_leader_clears_staging_partition_info() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.manifest_ctx.staging_partition_info().is_some()); + + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::Leader)); + assert!(region.manifest_ctx.staging_partition_info().is_none()); +} + +#[tokio::test] +async fn test_set_region_role_follower_clears_staging_partition_info() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.manifest_ctx.staging_partition_info().is_some()); + + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::Follower)); + assert!(region.manifest_ctx.staging_partition_info().is_none()); +} + +#[tokio::test] +async fn test_set_region_role_downgrading_leader_clears_staging_partition_info() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.manifest_ctx.staging_partition_info().is_some()); + + engine + .set_region_role(region_id, RegionRole::DowngradingLeader) + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::DowngradingLeader)); + assert!(region.manifest_ctx.staging_partition_info().is_none()); +} + +#[tokio::test] +async fn test_can_reenter_staging_after_direct_exit_cleanup() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(engine.role(region_id), Some(RegionRole::StagingLeader)); + assert!(region.manifest_ctx.staging_partition_info().is_some()); +} + async fn test_restricted_state_transitions_with_format(flat_format: bool) { let mut env = TestEnv::new().await; let engine = env diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index bd90779e0b..9846933d1f 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -547,7 +547,7 @@ async fn test_staging_manifest_directory_with_format(flat_format: bool) { .await .unwrap(); let region = engine.get_region(region_id).unwrap(); - let staging_partition_info = region.staging_partition_info.lock().unwrap().clone(); + let staging_partition_info = region.manifest_ctx.staging_partition_info(); assert_eq!( staging_partition_info .unwrap() diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 26ab96c779..3804b28afb 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -156,11 +156,6 @@ pub struct MitoRegion { pub(crate) topic_latest_entry_id: AtomicU64, /// The total bytes written to the region. pub(crate) written_bytes: Arc, - /// Partition info of the region in staging mode. - /// - /// During the staging mode, the region metadata in [`VersionControlRef`] is not updated, - /// so we need to store the partition info separately. - pub(crate) staging_partition_info: Mutex>, /// manifest stats stats: ManifestStats, } @@ -333,6 +328,17 @@ impl MitoRegion { self.manifest_ctx.set_role(next_role, self.region_id); } + pub(crate) fn region_role(&self) -> RegionRole { + match self.state() { + RegionRoleState::Follower => RegionRole::Follower, + RegionRoleState::Leader(RegionLeaderState::Staging) => RegionRole::StagingLeader, + RegionRoleState::Leader(RegionLeaderState::Downgrading) => { + RegionRole::DowngradingLeader + } + RegionRoleState::Leader(_) => RegionRole::Leader, + } + } + /// Sets the altering state. /// You should call this method in the worker loop. pub(crate) fn set_altering(&self) -> Result<()> { @@ -393,9 +399,8 @@ impl MitoRegion { /// You should call this method in the worker loop. /// Transitions from Staging to Writable state. pub fn exit_staging(&self) -> Result<()> { - *self.staging_partition_info.lock().unwrap() = None; - self.compare_exchange_state( - RegionLeaderState::Staging, + self.manifest_ctx.exit_staging( + self.region_id, RegionRoleState::Leader(RegionLeaderState::Writable), ) } @@ -819,7 +824,7 @@ impl MitoRegion { pub fn maybe_staging_partition_expr_str(&self) -> Option { let is_staging = self.is_staging(); if is_staging { - let staging_partition_info = self.staging_partition_info.lock().unwrap(); + let staging_partition_info = self.manifest_ctx.staging_partition_info(); if staging_partition_info.is_none() { warn!( "Staging partition expr is none for region {} in staging state", @@ -837,8 +842,8 @@ impl MitoRegion { pub fn expected_partition_expr_version(&self) -> u64 { if self.is_staging() { - let staging_partition_info = self.staging_partition_info.lock().unwrap(); - staging_partition_info + self.manifest_ctx + .staging_partition_info() .as_ref() .map(|info| info.partition_rule_version) .unwrap_or_default() @@ -852,8 +857,8 @@ impl MitoRegion { if !self.is_staging() { return false; } - let staging_partition_info = self.staging_partition_info.lock().unwrap(); - staging_partition_info + self.manifest_ctx + .staging_partition_info() .as_ref() .map(|info| { matches!( @@ -873,6 +878,11 @@ pub(crate) struct ManifestContext { /// The state of the region. The region checks the state before updating /// manifest. state: AtomicCell, + /// Partition info of the region in staging mode. + /// + /// During the staging mode, the region metadata in [`VersionControlRef`] is not updated, + /// so we need to store the partition info separately. + staging_partition_info: Mutex>, } impl ManifestContext { @@ -880,9 +890,46 @@ impl ManifestContext { ManifestContext { manifest_manager: tokio::sync::RwLock::new(manager), state: AtomicCell::new(state), + staging_partition_info: Mutex::new(None), } } + pub(crate) fn staging_partition_info(&self) -> Option { + self.staging_partition_info.lock().unwrap().clone() + } + + pub(crate) fn set_staging_partition_info(&self, staging_partition_info: StagingPartitionInfo) { + let mut current = self.staging_partition_info.lock().unwrap(); + debug_assert!(current.is_none()); + *current = Some(staging_partition_info); + } + + fn clear_staging_partition_info(&self) { + *self.staging_partition_info.lock().unwrap() = None; + } + + pub(crate) fn exit_staging( + &self, + region_id: RegionId, + next_state: RegionRoleState, + ) -> Result<()> { + self.state + .compare_exchange( + RegionRoleState::Leader(RegionLeaderState::Staging), + next_state, + ) + .map_err(|actual| { + RegionStateSnafu { + region_id, + state: actual, + expect: RegionRoleState::Leader(RegionLeaderState::Staging), + } + .build() + })?; + self.clear_staging_partition_info(); + Ok(()) + } + pub(crate) async fn manifest_version(&self) -> ManifestVersion { self.manifest_manager .read() @@ -1028,27 +1075,50 @@ impl ManifestContext { /// Sets the [`RegionRole`]. /// /// ```text - /// +------------------------------------------+ - /// | +-----------------+ | - /// | | | | - /// +---+------+ +-------+-----+ +--v-v---+ - /// | Follower | | Downgrading | | Leader | - /// +---^-^----+ +-----+-^-----+ +--+-+---+ - /// | | | | | | - /// | +------------------+ +-----------------+ | - /// +------------------------------------------+ - /// - /// Transition: - /// - Follower -> Leader - /// - Downgrading Leader -> Leader - /// - Leader -> Follower - /// - Downgrading Leader -> Follower - /// - Leader -> Downgrading Leader + /// +---------------------+ + /// | Staging Leader | + /// +----------+----------+ + /// | + /// v + /// +----------+ +------+-------+ +-------------+ + /// | Follower | <-> | Leader | <-> | Downgrading | + /// +-----+----+ +------+-------+ +------+------+ + /// ^ ^ | + /// +-----------------+--------------------+ /// /// ``` + /// + /// # State Transitions + /// + /// From `Follower`: + /// - `Follower -> Leader` + /// + /// From `Leader`: + /// - `Leader -> Follower` + /// - `Leader -> Downgrading Leader` + /// + /// From `Staging Leader`: + /// - `Staging Leader -> Leader` + /// - `Staging Leader -> Follower` + /// - `Staging Leader -> Downgrading Leader` + /// + /// From `Downgrading Leader`: + /// - `Downgrading Leader -> Leader` + /// - `Downgrading Leader -> Follower` pub(crate) fn set_role(&self, next_role: RegionRole, region_id: RegionId) { match next_role { RegionRole::Follower => { + if self + .exit_staging(region_id, RegionRoleState::Follower) + .is_ok() + { + info!( + "Convert region {} to follower, previous role state: {:?}", + region_id, + RegionRoleState::Leader(RegionLeaderState::Staging) + ); + return; + } match self.state.fetch_update(|state| { if !matches!(state, RegionRoleState::Follower) { Some(RegionRoleState::Follower) @@ -1071,6 +1141,20 @@ impl ManifestContext { } } RegionRole::Leader => { + if self + .exit_staging( + region_id, + RegionRoleState::Leader(RegionLeaderState::Writable), + ) + .is_ok() + { + info!( + "Convert region {} to leader, previous role state: {:?}", + region_id, + RegionRoleState::Leader(RegionLeaderState::Staging) + ); + return; + } match self.state.fetch_update(|state| { if matches!( state, @@ -1096,7 +1180,27 @@ impl ManifestContext { } } } + RegionRole::StagingLeader => { + info!( + "Ignore direct conversion of region {} to staging leader; staging requires the dedicated workflow", + region_id + ); + } RegionRole::DowngradingLeader => { + if self + .exit_staging( + region_id, + RegionRoleState::Leader(RegionLeaderState::Downgrading), + ) + .is_ok() + { + info!( + "Convert region {} to downgrading region, previous role state: {:?}", + region_id, + RegionRoleState::Leader(RegionLeaderState::Staging) + ); + return; + } match self.state.compare_exchange( RegionRoleState::Leader(RegionLeaderState::Writable), RegionRoleState::Leader(RegionLeaderState::Downgrading), @@ -1438,8 +1542,8 @@ pub fn parse_partition_expr(partition_expr_str: Option<&str>) -> Result StagingLeader should be ignored. + manifest_ctx.set_role(RegionRole::StagingLeader, region_id); + assert_eq!( + manifest_ctx.state.load(), + RegionRoleState::Leader(RegionLeaderState::Writable) + ); + // Leader -> Downgrading Leader manifest_ctx.set_role(RegionRole::DowngradingLeader, region_id); assert_eq!( @@ -1825,7 +1935,6 @@ mod tests { topic_latest_entry_id: Default::default(), written_bytes: Arc::new(AtomicU64::new(0)), stats: ManifestStats::default(), - staging_partition_info: Mutex::new(None), }; // Test initial state diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index b23e73557d..c1240c3829 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -17,7 +17,7 @@ use std::any::TypeId; use std::collections::HashMap; use std::sync::atomic::{AtomicI64, AtomicU64}; -use std::sync::{Arc, LazyLock, Mutex}; +use std::sync::{Arc, LazyLock}; use std::time::Instant; use common_telemetry::{debug, error, info, warn}; @@ -349,7 +349,6 @@ impl RegionOpener { topic_latest_entry_id: AtomicU64::new(0), written_bytes: Arc::new(AtomicU64::new(0)), stats: self.stats, - staging_partition_info: Mutex::new(None), })) } @@ -586,8 +585,6 @@ impl RegionOpener { topic_latest_entry_id: AtomicU64::new(topic_latest_entry_id), written_bytes: Arc::new(AtomicU64::new(0)), stats: self.stats.clone(), - // TODO(weny): reload the staging partition info from the manifest. - staging_partition_info: Mutex::new(None), }; let region = Arc::new(region); diff --git a/src/mito2/src/worker/handle_apply_staging.rs b/src/mito2/src/worker/handle_apply_staging.rs index e773150356..876d5c3c31 100644 --- a/src/mito2/src/worker/handle_apply_staging.rs +++ b/src/mito2/src/worker/handle_apply_staging.rs @@ -75,7 +75,7 @@ impl RegionWorkerLoop { return; } - let staging_partition_info = region.staging_partition_info.lock().unwrap().clone(); + let staging_partition_info = region.manifest_ctx.staging_partition_info(); let staging_partition_expr = staging_partition_info .as_ref() diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index 8b75fdd24f..83bd51df15 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -42,7 +42,7 @@ impl RegionWorkerLoop { // If the region is already in staging mode, verify the partition directive matches. if region.is_staging() { - let staging_partition_info = region.staging_partition_info.lock().unwrap().clone(); + let staging_partition_info = region.manifest_ctx.staging_partition_info(); // If the partition directive mismatches, return error. if staging_partition_info .as_ref() @@ -279,10 +279,8 @@ impl RegionWorkerLoop { region: &MitoRegionRef, partition_directive: StagingPartitionDirective, ) { - let mut staging_partition_info = region.staging_partition_info.lock().unwrap(); - debug_assert!(staging_partition_info.is_none()); - *staging_partition_info = Some(StagingPartitionInfo::from_partition_directive( - partition_directive, - )); + region.manifest_ctx.set_staging_partition_info( + StagingPartitionInfo::from_partition_directive(partition_directive), + ); } } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 287f64d225..b235fcffc7 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -67,7 +67,7 @@ impl From for RegionRole { SettableRegionRoleState::Follower => RegionRole::Follower, SettableRegionRoleState::DowngradingLeader => RegionRole::DowngradingLeader, SettableRegionRoleState::Leader => RegionRole::Leader, - SettableRegionRoleState::StagingLeader => RegionRole::Leader, // Still a leader role + SettableRegionRoleState::StagingLeader => RegionRole::StagingLeader, } } } @@ -210,6 +210,11 @@ pub enum RegionRole { Follower, // Writable region(mito2), Readonly region(file). Leader, + // Leader is in staging mode. + // + // This is leader-like and writable, but it follows the staging workflow + // semantics instead of a normal leader's steady state. + StagingLeader, // Leader is downgrading to follower. // // This state is used to prevent new write requests. @@ -221,6 +226,7 @@ impl Display for RegionRole { match self { RegionRole::Follower => write!(f, "Follower"), RegionRole::Leader => write!(f, "Leader"), + RegionRole::StagingLeader => write!(f, "Leader(Staging)"), RegionRole::DowngradingLeader => write!(f, "Leader(Downgrading)"), } } @@ -228,7 +234,7 @@ impl Display for RegionRole { impl RegionRole { pub fn writable(&self) -> bool { - matches!(self, RegionRole::Leader) + matches!(self, RegionRole::Leader | RegionRole::StagingLeader) } } @@ -237,6 +243,7 @@ impl From for PbRegionRole { match value { RegionRole::Follower => PbRegionRole::Follower, RegionRole::Leader => PbRegionRole::Leader, + RegionRole::StagingLeader => PbRegionRole::StagingLeader, RegionRole::DowngradingLeader => PbRegionRole::DowngradingLeader, } } @@ -246,6 +253,7 @@ impl From for RegionRole { fn from(value: PbRegionRole) -> Self { match value { PbRegionRole::Leader => RegionRole::Leader, + PbRegionRole::StagingLeader => RegionRole::StagingLeader, PbRegionRole::Follower => RegionRole::Follower, PbRegionRole::DowngradingLeader => RegionRole::DowngradingLeader, } From a24c58e25c5ab209d9034088073e894800e25a7b Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 13 Apr 2026 17:11:38 +0800 Subject: [PATCH 015/120] chore: fix git cliff errors in latest version (#7947) * chore: fix git cliff errors in latest version - Fix errors in v2.12.0 - Do not generate logs for beta/rc tags between the compared commits Signed-off-by: evenyag * chore: preserve blank line before release date in changelog Signed-off-by: evenyag --------- Signed-off-by: evenyag --- cliff.toml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cliff.toml b/cliff.toml index 4245203e92..2b35ddab5c 100644 --- a/cliff.toml +++ b/cliff.toml @@ -12,7 +12,9 @@ footer = "" body = """ # {{ version }} +{% if timestamp -%} Release date: {{ timestamp | date(format="%B %d, %Y") }} +{% endif -%} {%- set breakings = commits | filter(attribute="breaking", value=true) -%} {%- if breakings | length > 0 %} @@ -118,7 +120,10 @@ filter_commits = false # regex for skipping tags # skip_tags = "" # regex for ignoring tags -ignore_tags = ".*-nightly-.*" +# Ignore nightly tags and build-suffixed release tags such as +# v1.0.0-rc.2-13cdfa9b5-20260325-1774407105 so their commits are merged into +# the next visible release section instead of creating extra headings. +ignore_tags = ".*-nightly-.*|^v[0-9]+\\.[0-9]+\\.[0-9]+(-(alpha|beta|rc)\\.[0-9]+)?-[0-9a-f]{7,}-[0-9]{8}-[0-9]+$" # sort the tags topologically topo_order = false # sort the commits inside sections by oldest/newest order From 3750819f937d13408012ae36f621311c780fb7eb Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 13 Apr 2026 21:04:11 +0800 Subject: [PATCH 016/120] fix: match term zh (#7952) * fix: match term zh Signed-off-by: discord9 * chore: per gemini Signed-off-by: discord9 * chore: revert accident change Signed-off-by: discord9 * feat: unicode script han Signed-off-by: discord9 --------- Signed-off-by: discord9 --- Cargo.lock | 1 + Cargo.toml | 1 + src/common/function/Cargo.toml | 1 + .../function/src/scalars/matches_term.rs | 170 ++++++++++++++---- src/index/src/fulltext_index/tokenizer.rs | 20 +++ .../common/function/matches_term.result | 65 +++++++ .../common/function/matches_term.sql | 10 ++ 7 files changed, 238 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4f6339d83a..68d1dac297 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2348,6 +2348,7 @@ dependencies = [ "geohash", "h3o", "hyperloglogplus", + "icu_properties", "jsonb", "jsonpath-rust 0.7.5", "memchr", diff --git a/Cargo.toml b/Cargo.toml index 34e10d9173..66c35acee8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -161,6 +161,7 @@ humantime = "2.1" humantime-serde = "1.1" hyper = "1.1" hyper-util = "0.1" +icu_properties = "2.0.1" itertools = "0.14" jsonb = { version = "0.4.4", default-features = false } lazy_static = "1.4" diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index d164b9285d..43ddf9ae0c 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -47,6 +47,7 @@ geo-types = { version = "0.7", optional = true } geohash = { version = "0.13", optional = true } h3o = { version = "0.6", optional = true } hyperloglogplus = "0.4" +icu_properties.workspace = true jsonb.workspace = true jsonpath-rust = "0.7.5" memchr = "2.7" diff --git a/src/common/function/src/scalars/matches_term.rs b/src/common/function/src/scalars/matches_term.rs index 8dfb25cbc0..ec1b34d408 100644 --- a/src/common/function/src/scalars/matches_term.rs +++ b/src/common/function/src/scalars/matches_term.rs @@ -20,6 +20,8 @@ use datafusion_common::arrow::compute; use datafusion_common::arrow::datatypes::DataType; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; +use icu_properties::props::Script; +use icu_properties::{CodePointMapData, CodePointMapDataBorrowed}; use memchr::memmem; use crate::function::Function; @@ -27,10 +29,11 @@ use crate::function_registry::FunctionRegistry; /// Exact term/phrase matching function for text columns. /// -/// This function checks if a text column contains exact term/phrase matches -/// with non-alphanumeric boundaries. Designed for: -/// - Whole-word matching (e.g. "cat" in "cat!" but not in "category") +/// This function uses script-aware matching rules: +/// - ASCII-only terms keep whole-word style boundary matching, like Whole-word matching (e.g. "cat" in "cat!" but not in "category") /// - Phrase matching (e.g. "hello world" in "note:hello world!") +/// - Terms containing Han characters match as contiguous substrings +/// - Mixed-script identifiers and numeric terms remain searchable in Chinese text /// /// # Signature /// `matches_term(text: String, term: String) -> Boolean` @@ -43,9 +46,8 @@ use crate::function_registry::FunctionRegistry; /// BooleanVector where each element indicates if the corresponding text /// contains an exact match of the term, following these rules: /// 1. Exact substring match found (case-sensitive) -/// 2. Match boundaries are either: -/// - Start/end of text -/// - Any non-alphanumeric character (including spaces, hyphens, punctuation, etc.) +/// 2. For ASCII-only terms, adjacent ASCII word characters block the match +/// 3. For Han-containing terms, contiguous substring match is sufficient /// /// # Examples /// ``` @@ -60,6 +62,9 @@ use crate::function_registry::FunctionRegistry; /// SELECT matches_term(column, 'critical error') FROM logs; /// -- Match in: "ERROR:critical error!" /// -- No match: "critical_errors" +/// -- Chinese substring examples -- +/// SELECT matches_term(column, '手机') FROM table; +/// -- Text: "登录手机号18888888888的动态key" => true /// /// -- Empty string handling -- /// SELECT matches_term(column, '') FROM table; @@ -204,9 +209,8 @@ impl Function for MatchesTermFunction { /// /// A term is considered matched when: /// 1. The exact sequence appears in the text -/// 2. It is either: -/// - At the start/end of text with adjacent non-alphanumeric character -/// - Surrounded by non-alphanumeric characters +/// 2. ASCII-only terms are not adjacent to ASCII word characters +/// 3. Han-containing terms match as contiguous substrings /// /// # Examples /// ``` @@ -215,28 +219,105 @@ impl Function for MatchesTermFunction { /// assert!(finder.find("dog,cat")); // Term preceded by comma /// assert!(!finder.find("category")); // Partial match rejected /// -/// let finder = MatchesTermFinder::new("world"); -/// assert!(finder.find("hello-world")); // Hyphen boundary +/// let finder = MatchesTermFinder::new("手机"); +/// assert!(finder.find("登录手机号18888888888的动态key")); /// ``` #[derive(Clone, Debug)] pub struct MatchesTermFinder { finder: memmem::Finder<'static>, term: String, - starts_with_non_alnum: bool, - ends_with_non_alnum: bool, + term_kind: TermKind, + starts_with_other: bool, + ends_with_other: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum CharClass { + AsciiWord, + Han, + UnicodeWord, + Other, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TermKind { + AsciiLike, + UnicodeWord, + HanContaining, +} + +fn classify_char(c: char) -> CharClass { + if c.is_ascii_alphanumeric() { + CharClass::AsciiWord + } else if is_han(c) { + CharClass::Han + } else if c.is_alphanumeric() { + CharClass::UnicodeWord + } else { + CharClass::Other + } +} + +static HAN_SCRIPT_DATA: CodePointMapDataBorrowed<'static, Script> = + CodePointMapData::