diff --git a/src/common/meta/src/ddl/create_table/executor.rs b/src/common/meta/src/ddl/create_table/executor.rs index 88d88fe277c..332c9e57eb4 100644 --- a/src/common/meta/src/ddl/create_table/executor.rs +++ b/src/common/meta/src/ddl/create_table/executor.rs @@ -182,7 +182,19 @@ impl CreateTableExecutor { region_wal_options: RegionWalOptions, ) -> Result<()> { if !column_metadatas.is_empty() { - update_table_info_column_ids(&mut table_info, column_metadatas); + // Region metadata can include engine-internal columns that are not visible in + // the table schema. Only visible columns belong in the table's column IDs. + let visible_column_metadatas = column_metadatas + .iter() + .filter(|column| { + table_info + .meta + .schema + .contains_column(&column.column_schema.name) + }) + .cloned() + .collect::>(); + update_table_info_column_ids(&mut table_info, &visible_column_metadatas); } let detecting_regions = convert_region_routes_to_detecting_regions(&table_route.region_routes); diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index 27950f0fd5d..8defc6bd0f2 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -25,8 +25,8 @@ use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ - ResolveColumnMetadataResult, build_column_metadata_from_table_info, - check_column_metadatas_consistent, resolve_column_metadatas_with_latest, + ResolveColumnMetadataResult, build_reconciliation_column_metadata, + check_column_metadatas_consistent, reorder_tag_columns, resolve_column_metadatas_with_latest, resolve_column_metadatas_with_metasrv, }; @@ -92,6 +92,8 @@ impl State for ResolveColumnMetadata { ctx.persistent_ctx.table_info_value = Some(table_info_value); if let Some(column_metadatas) = check_column_metadatas_consistent(&self.region_metadata) { + let column_metadatas = + reorder_tag_columns(&column_metadatas, &self.region_metadata[0].primary_key)?; // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); info!( @@ -125,7 +127,7 @@ impl State for ResolveColumnMetadata { .table_info .name_to_ids() .context(MissingColumnIdsSnafu)?; - let column_metadata = build_column_metadata_from_table_info( + let column_metadata = build_reconciliation_column_metadata( table_info_value.table_info.meta.schema.column_schemas(), &table_info_value.table_info.meta.primary_key_indices, &name_to_ids, diff --git a/src/common/meta/src/reconciliation/utils.rs b/src/common/meta/src/reconciliation/utils.rs index 6ddc084596b..326516ecd80 100644 --- a/src/common/meta/src/reconciliation/utils.rs +++ b/src/common/meta/src/reconciliation/utils.rs @@ -85,6 +85,70 @@ pub(crate) fn check_column_metadatas_consistent( Some(region_metadatas[0].column_metadatas.clone()) } +/// Returns columns with tag values in primary-key order while preserving the +/// positions of non-tag columns. +pub(crate) fn reorder_tag_columns( + column_metadatas: &[ColumnMetadata], + primary_key: &[u32], +) -> Result> { + let tag_count = column_metadatas + .iter() + .filter(|column| column.semantic_type == SemanticType::Tag) + .count(); + ensure!( + primary_key.len() == tag_count, + UnexpectedSnafu { + err_msg: format!( + "Number of primary key columns {} does not match tag columns {}", + primary_key.len(), + tag_count, + ), + } + ); + + let mut primary_key_ids = HashSet::with_capacity(primary_key.len()); + let mut tags = Vec::with_capacity(primary_key.len()); + for column_id in primary_key { + let column = column_metadatas + .iter() + .find(|column| column.column_id == *column_id) + .with_context(|| UnexpectedSnafu { + err_msg: format!( + "Primary key column {} not found in column metadata", + column_id + ), + })?; + ensure!( + column.semantic_type == SemanticType::Tag, + UnexpectedSnafu { + err_msg: format!("Primary key column {} is not a tag", column_id), + } + ); + ensure!( + primary_key_ids.insert(column_id), + UnexpectedSnafu { + err_msg: format!("Primary key column {} is duplicated", column_id), + } + ); + tags.push(column); + } + + let mut tags = tags.into_iter(); + + column_metadatas + .iter() + .map(|column| { + if column.semantic_type == SemanticType::Tag { + tags.next().cloned().with_context(|| UnexpectedSnafu { + err_msg: "Primary key has fewer columns than tags".to_string(), + }) + } else { + Ok(column.clone()) + } + }) + .collect() +} + /// Resolves column metadata inconsistencies among the given region metadatas /// by using the column metadata from the metasrv as the source of truth. /// @@ -157,8 +221,13 @@ pub(crate) fn resolve_column_metadatas_with_latest( } } - // TODO(weny): verify the new column metadatas are acceptable for regions. - Ok((latest_region_metadata.column_metadatas.clone(), region_ids)) + Ok(( + reorder_tag_columns( + &latest_region_metadata.column_metadatas, + &latest_region_metadata.primary_key, + )?, + region_ids, + )) } /// Constructs a vector of [`ColumnMetadata`] from the provided table information. @@ -206,6 +275,22 @@ pub(crate) fn build_column_metadata_from_table_info( .collect::>>() } +/// Builds SyncColumns metadata with tags in the TableInfo primary-key order. +/// Non-tag columns retain their schema positions. +pub(crate) fn build_reconciliation_column_metadata( + column_schemas: &[ColumnSchema], + primary_key_indexes: &[usize], + name_to_ids: &HashMap, +) -> Result> { + let column_metadatas = + build_column_metadata_from_table_info(column_schemas, primary_key_indexes, name_to_ids)?; + let primary_key = primary_key_indexes + .iter() + .map(|index| column_metadatas[*index].column_id) + .collect::>(); + reorder_tag_columns(&column_metadatas, &primary_key) +} + /// Checks whether the schema invariants hold between the existing and new column metadata. /// /// Invariants: @@ -1180,6 +1265,104 @@ mod tests { assert_matches!(err, Error::MissingColumnInColumnMetadata { .. }); } + #[test] + fn test_reconciliation_columns_follow_primary_key_order() { + let columns = vec![ + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_b", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("field", ConcreteDataType::int32_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 2, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_a", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 3, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + semantic_type: SemanticType::Timestamp, + column_id: 4, + }, + ]; + let schemas = columns + .iter() + .map(|column| column.column_schema.clone()) + .collect::>(); + let ids = columns + .iter() + .map(|column| (column.column_schema.name.clone(), column.column_id)) + .collect(); + let expected = vec![3, 2, 1, 4]; + assert_eq!( + build_reconciliation_column_metadata(&schemas, &[2, 0], &ids) + .unwrap() + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + + let mut metadata = build_region_metadata(RegionId::new(1024, 0), &columns); + metadata.primary_key = vec![3, 1]; + metadata.schema_version = 2; + + assert_eq!( + check_column_metadatas_consistent(&[metadata.clone()]).unwrap(), + columns + ); + + assert_eq!( + resolve_column_metadatas_with_latest(&[metadata]) + .unwrap() + .0 + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + } + + #[test] + fn test_reorder_tag_columns_rejects_invalid_primary_key() { + let columns = new_test_column_metadatas(); + + let err = reorder_tag_columns(&columns, &[999]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("not found in column metadata")); + + let err = reorder_tag_columns(&columns, &[]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("does not match tag columns")); + + let err = reorder_tag_columns(&columns, &[2]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is not a tag")); + + let mut duplicate_tags = columns.clone(); + duplicate_tags[2].semantic_type = SemanticType::Tag; + let err = reorder_tag_columns(&duplicate_tags, &[0, 0]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is duplicated")); + } + #[test] fn test_check_column_metadatas_consistent() { let column_metadatas = new_test_column_metadatas(); diff --git a/src/meta-srv/src/handler/failure_handler.rs b/src/meta-srv/src/handler/failure_handler.rs index fa089c14c2a..07484ac9b35 100644 --- a/src/meta-srv/src/handler/failure_handler.rs +++ b/src/meta-srv/src/handler/failure_handler.rs @@ -52,9 +52,10 @@ impl HeartbeatHandler for RegionFailureHandler { return Ok(HandleControl::Continue); }; + // Never blocks: the region lease has already been decided earlier in the + // chain, so a busy supervisor must not withhold the heartbeat response. self.heartbeat_acceptor - .accept(DatanodeHeartbeat::from(stat)) - .await; + .accept(DatanodeHeartbeat::from(stat)); Ok(HandleControl::Continue) } @@ -62,6 +63,8 @@ impl HeartbeatHandler for RegionFailureHandler { #[cfg(test)] mod tests { + use std::time::Duration; + use api::v1::meta::HeartbeatRequest; use common_catalog::consts::default_engine; use common_meta::datanode::{RegionManifestInfo, RegionStat, Stat}; @@ -73,7 +76,43 @@ mod tests { use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::builder::MetasrvBuilder; use crate::region::supervisor::tests::new_test_supervisor; - use crate::region::supervisor::{Event, HeartbeatAcceptor}; + use crate::region::supervisor::{DatanodeHeartbeat, Event, HeartbeatAcceptor}; + + fn new_region_stat(region_id: u64) -> RegionStat { + RegionStat { + id: RegionId::from_u64(region_id), + rcus: 0, + wcus: 0, + approximate_bytes: 0, + engine: default_engine().to_string(), + role: RegionRole::Follower, + num_rows: 0, + memtable_size: 0, + manifest_size: 0, + sst_size: 0, + sst_num: 0, + index_size: 0, + region_manifest: RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + } + } + + fn new_stat() -> Stat { + Stat { + id: 42, + region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], + timestamp_millis: 1000, + ..Default::default() + } + } #[tokio::test] async fn test_handle_heartbeat() { @@ -85,38 +124,7 @@ mod tests { let metasrv = builder.build().await.unwrap(); let mut ctx = metasrv.new_ctx(); let acc = &mut HeartbeatAccumulator::default(); - fn new_region_stat(region_id: u64) -> RegionStat { - RegionStat { - id: RegionId::from_u64(region_id), - rcus: 0, - wcus: 0, - approximate_bytes: 0, - engine: default_engine().to_string(), - role: RegionRole::Follower, - num_rows: 0, - memtable_size: 0, - manifest_size: 0, - sst_size: 0, - sst_num: 0, - index_size: 0, - region_manifest: RegionManifestInfo::Mito { - manifest_version: 0, - flushed_entry_id: 0, - file_removed_cnt: 0, - }, - data_topic_latest_entry_id: 0, - metadata_topic_latest_entry_id: 0, - written_bytes: 0, - query_cpu_time: 0, - query_scanned_bytes: 0, - } - } - acc.stat = Some(Stat { - id: 42, - region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], - timestamp_millis: 1000, - ..Default::default() - }); + acc.stat = Some(new_stat()); handler.handle(req, &mut ctx, acc).await.unwrap(); let (tx, rx) = oneshot::channel(); @@ -124,4 +132,39 @@ mod tests { let detector = rx.await.unwrap(); assert_eq!(detector.iter().collect::>().len(), 3); } + + #[tokio::test] + async fn test_handle_heartbeat_does_not_block_on_full_supervisor_queue() { + // Regression test for #9419: a full queue used to park the handler. + // Keep the supervisor alive, or the queue reports Closed. + let (_supervisor, sender) = new_test_supervisor(); + // Built directly, so the supervisor stays undrained. + let handler = RegionFailureHandler { + heartbeat_acceptor: HeartbeatAcceptor::new(sender.clone()), + }; + + let stat = new_stat(); + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(DatanodeHeartbeat::from(&stat))) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + + let builder = MetasrvBuilder::new(); + let metasrv = builder.build().await.unwrap(); + let mut ctx = metasrv.new_ctx(); + let acc = &mut HeartbeatAccumulator::default(); + // Without a stat the handler returns before the acceptor. + acc.stat = Some(stat); + + let result = tokio::time::timeout( + Duration::from_secs(5), + handler.handle(&HeartbeatRequest::default(), &mut ctx, acc), + ) + .await; + assert!(result.is_ok()); + } } diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 1cf32194302..c92e6f31fff 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -48,6 +48,9 @@ lazy_static! { /// Heartbeat received by metasrv. pub static ref METRIC_META_HEARTBEAT_RECV: IntCounterVec = register_int_counter_vec!("greptime_meta_heartbeat_recv", "heartbeats received by metasrv", &["pusher_key"]).unwrap(); + /// Heartbeats dropped by metasrv. + pub static ref METRIC_META_HEARTBEAT_DROPPED: IntCounterVec = + register_int_counter_vec!("greptime_meta_heartbeat_dropped", "heartbeats dropped by metasrv", &["reason"]).unwrap(); /// The migration execute histogram. pub static ref METRIC_META_REGION_MIGRATION_EXECUTE: HistogramVec = register_histogram_vec!("greptime_meta_region_migration_execute", "meta region migration execute", &["state"]).unwrap(); diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 56a9fe4ee57..99c68ac23c8 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -35,6 +35,7 @@ use common_time::util::current_time_millis; use futures::{StreamExt, TryStreamExt}; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; +use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::oneshot; use tokio::time::{MissedTickBehavior, interval, interval_at}; @@ -43,6 +44,7 @@ use crate::discovery::utils::accept_ingest_workload; use crate::error::{self, Result}; use crate::failure_detector::PhiAccrualFailureDetectorOptions; use crate::metasrv::{RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef}; +use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::manager::{ RegionMigrationManagerRef, RegionMigrationTriggerReason, SubmitRegionMigrationTaskResult, }; @@ -356,10 +358,29 @@ impl HeartbeatAcceptor { Self { sender } } - /// Accepts heartbeats from datanodes. - pub(crate) async fn accept(&self, heartbeat: DatanodeHeartbeat) { - if let Err(err) = self.sender.send(Event::HeartbeatArrived(heartbeat)).await { - error!(err; "RegionSupervisor has stop receiving heartbeat."); + /// Accepts heartbeats from datanodes without waiting for queue capacity: a full queue + /// drops the heartbeat, counted in [`METRIC_META_HEARTBEAT_DROPPED`]. A dropped + /// heartbeat costs the failure detector one sample. + pub(crate) fn accept(&self, heartbeat: DatanodeHeartbeat) { + let datanode_id = heartbeat.datanode_id; + match self.sender.try_send(Event::HeartbeatArrived(heartbeat)) { + Ok(()) => {} + Err(TrySendError::Full(_)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .inc(); + warn!( + "Dropping heartbeat because the region supervisor event queue is full, datanode_id: {}, queue_capacity: {}", + datanode_id, + self.sender.max_capacity() + ); + } + Err(TrySendError::Closed(err)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["closed"]) + .inc(); + error!(err; "RegionSupervisor has stop receiving heartbeat."); + } } } } @@ -899,14 +920,15 @@ pub(crate) mod tests { use tokio::time::sleep; use super::RegionSupervisorSelector; + use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::RegionMigrationTriggerReason; use crate::procedure::region_migration::manager::{ RegionMigrationManager, SubmitRegionMigrationTaskResult, }; use crate::procedure::region_migration::test_util::TestingEnv; use crate::region::supervisor::{ - DatanodeHeartbeat, Event, RegionFailureDetectorControl, RegionSupervisor, - RegionSupervisorTicker, + DatanodeHeartbeat, Event, HeartbeatAcceptor, RegionFailureDetectorControl, + RegionSupervisor, RegionSupervisorTicker, }; use crate::selector::test_utils::{RandomNodeSelector, new_test_selector_context}; @@ -940,6 +962,14 @@ pub(crate) mod tests { ) } + fn filler_heartbeat() -> DatanodeHeartbeat { + DatanodeHeartbeat { + datanode_id: 1, + regions: vec![], + timestamp: 0, + } + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); @@ -1372,4 +1402,39 @@ pub(crate) mod tests { assert!(supervisor.failure_detector.contains(&detecting_region)); assert!(supervisor.failover_counts.contains_key(&detecting_region)); } + + #[tokio::test] + async fn test_heartbeat_accept_does_not_block_on_full_queue() { + // Regression test for #9419: a full queue used to park the caller. + // The supervisor is never run, so the queue stays full. + let (mut supervisor, sender) = new_test_supervisor(); + + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(filler_heartbeat())) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + assert_eq!(sender.capacity(), 0); + + let acceptor = HeartbeatAcceptor::new(sender.clone()); + let dropped = || { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .get() + }; + let dropped_before = dropped(); + // The full queue drops the heartbeat and counts it. + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + + // Free one slot: the next heartbeat is accepted. + supervisor.receiver.try_recv().unwrap(); + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + } } diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index fd74c4b498a..cf039b52278 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -37,7 +37,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::provider::Provider; -use store_api::metadata::ColumnMetadata; +use store_api::metadata::{ColumnMetadata, RegionMetadata}; use store_api::metric_engine_consts::{PRIMARY_KEY_ENCODING, TABLE_COLUMN_METADATA_EXTENSION_KEY}; use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ @@ -193,6 +193,133 @@ fn assert_column_metadatas(column_name: &[(&str, ColumnId)], column_metadatas: & } } +#[tokio::test] +async fn test_sync_columns_preserves_primary_key_order() { + test_sync_columns_preserves_primary_key_order_with_format(false).await; + test_sync_columns_preserves_primary_key_order_with_format(true).await; +} + +async fn assert_sync_columns_readable( + engine: &MitoEngine, + region_id: RegionId, + metadata: &RegionMetadata, + expected: &str, +) { + let current = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(current.primary_key, metadata.primary_key); + let projection = metadata + .column_metadatas + .iter() + .map(|column| current.column_index_by_id(column.column_id).unwrap()) + .collect(); + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + projection: Some(projection), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap(), + expected + ); +} + +async fn test_sync_columns_preserves_primary_key_order_with_format(flat_format: bool) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let mut row_schema = rows_schema(&request); + row_schema.push(tag_column_schema("tag_1", ColumnDataType::String)); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + engine + .handle_request(region_id, RegionRequest::Alter(add_tag1())) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(metadata.primary_key, vec![0, 3]); + + put_rows( + &engine, + region_id, + Rows { + schema: row_schema, + rows: build_rows_for_tags("a", "b", 0, 2, 10), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let before = RecordBatches::try_collect( + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .scan() + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + let mut columns = metadata.column_metadatas.clone(); + columns.push(ColumnMetadata { + column_schema: ColumnSchema::new("field_1", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 4, + }); + let tags = metadata.primary_key_columns().cloned().collect::>(); + let mut tags = tags.into_iter(); + for column in &mut columns { + if column.semantic_type == SemanticType::Tag { + *column = tags.next().unwrap(); + } + } + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SyncColumns { + column_metadatas: columns, + }, + }), + ) + .await + .unwrap(); + + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; +} + #[tokio::test] async fn test_alter_region() { test_alter_region_with_format(false).await; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 9f1051f3282..59c80c6dfbc 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -71,7 +71,7 @@ use crate::sst::SeriesEstimator; use crate::sst::index::IndexOutput; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; -use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo}; +use crate::sst::parquet::{COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo}; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1348,7 +1348,7 @@ impl BulkPartEncoder { .set_write_batch_size(row_group_size) .set_max_row_group_row_count(Some(row_group_size)) .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_column_index_truncate_length(None) + .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None) .build(), ); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 31ca70442b1..b980bf1c30a 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -72,6 +72,18 @@ pub(crate) struct Json2TargetLayout { /// batching without changing the row group layout of newly written SSTs. pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; +/// Truncation length of min/max values in the parquet column index. +/// +/// Mito never reads the column index (it is skipped or stripped on load), so truncating +/// it is safe. Chunk statistics stay untruncated because pruning decodes primary keys +/// from them. Without truncation every page of a large string column stores its whole +/// min and max values, uncompressed. +/// +/// Truncated values are only bounds: code that starts reading the column index must not +/// decode primary keys from it. +pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = + parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; + /// Parquet write options. #[derive(Debug, Clone)] pub struct WriteOptions { @@ -144,6 +156,7 @@ mod tests { use datatypes::schema::{FulltextAnalyzer, FulltextBackend, FulltextOptions}; use object_store::ObjectStore; use parquet::arrow::AsyncArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; @@ -697,6 +710,84 @@ mod tests { .await; } + #[tokio::test] + async fn test_column_index_truncates_large_field_values() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let handle = sst_file_handle(0, 1000); + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let metadata = build_test_binary_test_region_metadata(); + // A tag long enough that the encoded primary key exceeds the 64-byte truncation length. + let tag = "t".repeat(100); + // parquet 58 checks the page size limit once per write batch (1024 rows), so + // write enough rows to flush several data pages. + let values: Vec> = (0..3000) + .map(|i| format!("{i:08}").into_bytes().repeat(128)) + .collect(); + let batch = new_record_batch_with_binary_values(&tag, &values); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + writer + .write_all_flat_as_primary_key( + new_flat_source_from_record_batches(vec![batch]), + None, + &WriteOptions::default(), + ) + .await + .unwrap(); + + let path = handle.file_path(FILE_DIR, PathType::Bare); + let bytes = object_store.read(&path).await.unwrap().to_bytes(); + let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap(); + let parquet_meta = builder.metadata().clone(); + let schema = parquet_meta.file_metadata().schema_descr(); + let column = |name: &str| { + (0..schema.num_columns()) + .find(|i| schema.column(*i).name() == name) + .unwrap() + }; + + let field_index = &parquet_meta.column_index().unwrap()[0][column("field_0")]; + let parquet::file::page_index::column_index::ColumnIndexMetaData::BYTE_ARRAY(field_index) = + field_index + else { + panic!("unexpected column index: {field_index:?}"); + }; + let mut pages = 0; + for (min, max) in field_index + .min_values_iter() + .zip(field_index.max_values_iter()) + { + assert!(min.unwrap().len() <= 64); + assert!(max.unwrap().len() <= 64); + pages += 1; + } + assert!(pages > 1, "values should span several pages"); + + // Chunk statistics keep exact primary keys; pruning decodes them. + let pk_stats = parquet_meta + .row_group(0) + .column(column(store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME)) + .statistics() + .unwrap(); + let pk = new_primary_key(&[&tag]); + assert!(pk.len() > 64); + assert_eq!(pk_stats.max_bytes_opt().unwrap(), pk.as_slice()); + } + #[rstest::rstest] #[tokio::test] async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { @@ -1272,6 +1363,32 @@ mod tests { assert!(cached.contains_row_group(3)); } + fn new_record_batch_with_binary_values(tag: &str, values: &[Vec]) -> RecordBatch { + let metadata = build_test_binary_test_region_metadata(); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + let num_rows = values.len(); + let mut tag_0_builder = StringDictionaryBuilder::::new(); + let mut pk_builder = BinaryDictionaryBuilder::::new(); + let pk = new_primary_key(&[tag]); + for _ in 0..num_rows { + tag_0_builder.append_value(tag); + pk_builder.append(&pk).unwrap(); + } + let columns: Vec = vec![ + Arc::new(tag_0_builder.finish()), + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values.iter().map(|v| v.as_slice()), + )), + Arc::new(TimestampMillisecondArray::from_iter_values( + 0..num_rows as i64, + )), + Arc::new(pk_builder.finish()), + Arc::new(UInt64Array::from_value(1000, num_rows)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)), + ]; + RecordBatch::try_new(flat_schema, columns).unwrap() + } + fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch { assert!(end >= start); let metadata = build_test_binary_test_region_metadata(); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 45c8bdcbbe6..b886214882f 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -61,7 +61,9 @@ use crate::sst::parquet::flat_format::{ FlatWriteFormat, primary_key_column_index, time_index_column_index, }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; -use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; +use crate::sst::parquet::{ + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, +}; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, }; @@ -565,7 +567,7 @@ where .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_encoding(Encoding::PLAIN) .set_max_row_group_row_count(Some(opts.row_group_size)) - .set_column_index_truncate_length(None) + .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); let props_builder = Self::customize_column_config(props_builder, &self.metadata); diff --git a/src/operator/src/tests/partition_manager.rs b/src/operator/src/tests/partition_manager.rs index 7fdcb96d93f..b5423050780 100644 --- a/src/operator/src/tests/partition_manager.rs +++ b/src/operator/src/tests/partition_manager.rs @@ -366,3 +366,127 @@ async fn test_logical_split_rows_rejects_physical_only_partition_column() { partition::error::Error::UndefinedColumn { ref column, .. } if column == "b" )); } + +#[tokio::test] +async fn test_partition_routing_with_empty_default_expression() { + use datatypes::arrow::array::{ArrayRef, RecordBatch, StringArray}; + use datatypes::value::Value; + use partition::expr::col; + use store_api::storage::RegionId; + + for (transition, target_count) in [(true, 1), (true, 3), (false, 3)] { + let backend = Arc::new(common_meta::kv_backend::memory::MemoryKvBackend::new()); + let route_cache = test_new_table_route_cache(backend.clone()); + let manager = PartitionRuleManager::new( + backend.clone(), + route_cache.clone(), + test_new_partition_info_cache(route_cache), + ); + let value = |s: &str| Value::String(s.into()); + let mut exprs = vec![ + col("host").lt(value("36z0")), + col("host") + .gt_eq(value("36z0")) + .and(col("host").lt(value("7Ez0"))), + col("host") + .gt_eq(value("7Ez0")) + .and(col("host").lt(value("BMz0"))), + col("host").gt_eq(value("BMz0")), + ]; + if target_count == 1 { + exprs = vec![ + col("host").lt(value("36z0")), + col("host").gt_eq(value("36z0")), + ]; + } + let routes = exprs + .into_iter() + .enumerate() + .map(|(i, expr)| RegionRoute { + region: Region { + id: RegionId::new(1027, i as u32), + partition_expr: if transition && i == 0 { + String::new() + } else { + expr.as_json_str().unwrap() + }, + ..Default::default() + }, + leader_peer: Some(Peer::new(1, "")), + ..Default::default() + }) + .collect(); + let info = new_test_table_info_with_columns( + 1027, + "physical", + vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ], + vec![0], + ); + let metadata_manager = TableMetadataManager::new(backend); + metadata_manager + .create_table_metadata( + info.clone(), + TableRouteValue::physical(routes), + new_test_region_wal_options((0..=target_count).collect()), + ) + .await + .unwrap(); + let mut logical_info = info.clone(); + logical_info.ident.table_id = 1028; + logical_info.name = "logical".into(); + metadata_manager + .create_table_metadata( + logical_info.clone(), + TableRouteValue::logical(1027), + HashMap::new(), + ) + .await + .unwrap(); + + let hosts = [None, Some("1"), Some("5"), Some("8"), Some("animi")]; + let expected_regions = if target_count == 1 { + [0, 0, 1, 1, 1] + } else { + [0, 0, 1, 2, 3] + }; + let batch = RecordBatch::try_from_iter([( + "host", + Arc::new(StringArray::from(hosts.to_vec())) as ArrayRef, + )]) + .unwrap(); + for table_info in [&info, &logical_info] { + let (rule, _) = manager.find_table_partition_rule(table_info).await.unwrap(); + for (host, expected) in hosts.into_iter().zip(expected_regions) { + assert_eq!( + rule.find_region(&[host.map(value).unwrap_or(Value::Null)]) + .unwrap(), + expected + ); + } + let splits = rule.split_record_batch(&batch).unwrap(); + assert_eq!(splits.len(), target_count as usize + 1); + for (region, mask) in splits { + let selected = mask + .array() + .iter() + .enumerate() + .filter_map(|(i, selected)| (selected == Some(true)).then_some(i)) + .collect::>(); + let expected = expected_regions + .iter() + .enumerate() + .filter_map(|(i, expected)| (*expected == region).then_some(i)) + .collect::>(); + assert_eq!(selected, expected); + } + } + } +} diff --git a/src/partition/src/manager.rs b/src/partition/src/manager.rs index 89bb61c4174..68c9586e908 100644 --- a/src/partition/src/manager.rs +++ b/src/partition/src/manager.rs @@ -206,17 +206,25 @@ impl PartitionRuleManager { .iter() .map(|r| (r.id.region_number(), r.partition_expr_version)) .collect::>>(); - let regions = partition_info + let (mut regions, exprs): (Vec<_>, Vec<_>) = partition_info .partitions .iter() - .map(|x| x.id.region_number()) - .collect::>(); - let exprs = partition_info - .partitions - .iter() - .filter_map(|x| x.partition_expr.as_ref()) - .cloned() - .collect::>(); + .filter_map(|partition| { + partition + .partition_expr + .as_ref() + .map(|expr| (partition.id.region_number(), expr.clone())) + }) + .unzip(); + // Empty expressions belong to the default region. Keep it after the + // explicit regions so expression indices still identify the right region. + regions.extend( + partition_info + .partitions + .iter() + .filter(|partition| partition.partition_expr.is_none()) + .map(|partition| partition.id.region_number()), + ); let partition_rule = Arc::new(MultiDimPartitionRule::try_new( partition_columns, regions, diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index aae9df85e40..cf9e8b0b59e 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -384,10 +384,7 @@ impl HistogramFold { let output_schema: SchemaRef = self.output_schema.inner().clone(); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), - Partitioning::Hash( - partition_exprs.clone(), - exec_input.output_partitioning().partition_count(), - ), + Partitioning::UnknownPartitioning(exec_input.output_partitioning().partition_count()), EmissionType::Incremental, Boundedness::Bounded, )); @@ -595,10 +592,7 @@ impl ExecutionPlan for HistogramFoldExec { let new_input = children[0].clone(); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(self.output_schema.clone()), - Partitioning::Hash( - self.partition_exprs.clone(), - new_input.output_partitioning().partition_count(), - ), + Partitioning::UnknownPartitioning(new_input.output_partitioning().partition_count()), EmissionType::Incremental, Boundedness::Bounded, )); @@ -1448,7 +1442,9 @@ mod test { RecordBatch::try_new(schema, columns).unwrap() } - fn prepare_test_data() -> DataSourceExec { + /// The batches and schema shared by [`prepare_test_data`] and + /// [`prepare_test_data_with_two_partitions`]. + fn prepare_test_data_batches() -> (Vec, SchemaRef) { let schema = Arc::new(Schema::new(vec![ Field::new("host", DataType::Utf8, true), Field::new("le", DataType::Utf8, true), @@ -1501,8 +1497,24 @@ mod test { ) .unwrap(); + (vec![data_1, data_2, data_3], schema) + } + + fn prepare_test_data() -> DataSourceExec { + let (batches, schema) = prepare_test_data_batches(); DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![data_1, data_2, data_3]], schema, None).unwrap(), + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + )) + } + + /// The batches from [`prepare_test_data_batches`], laid out over two + /// partitions, so that partitioning metadata is checked against a real + /// multi-partition input. + fn prepare_test_data_with_two_partitions() -> DataSourceExec { + let (mut batches, schema) = prepare_test_data_batches(); + let first_partition = vec![batches.remove(0)]; + DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[first_partition, batches], schema, None).unwrap(), )) } @@ -1590,9 +1602,8 @@ mod test { let properties = PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), - Partitioning::Hash( - partition_exprs.clone(), - input.output_partitioning().partition_count(), + Partitioning::UnknownPartitioning( + input.properties().output_partitioning().partition_count(), ), EmissionType::Incremental, Boundedness::Bounded, @@ -1795,6 +1806,34 @@ mod test { assert!(broken_result.is_err()); } + #[test] + fn fold_does_not_claim_hash_partitioning() { + let input: Arc = Arc::new(prepare_test_data_with_two_partitions()); + let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(input.schema().to_dfschema().unwrap()), + }); + let logical_fold = HistogramFold::new( + "le".to_string(), + "val".to_string(), + "host".to_string(), + 0.5, + logical_input, + ) + .unwrap(); + let constructed = logical_fold.to_execution_plan(input.clone()); + assert!(matches!( + constructed.output_partitioning(), + Partitioning::UnknownPartitioning(2) + )); + + let rebuilt = constructed.with_new_children(vec![input]).unwrap(); + assert!(matches!( + rebuilt.output_partitioning(), + Partitioning::UnknownPartitioning(2) + )); + } + #[test] fn confirm_schema() { let input_schema = Schema::new(vec![ diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 6b7bbd3fab1..3ff0a14f411 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -348,12 +348,13 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } -/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +/// Non-scalar values (arrays, maps, bytes) are not representable as tags. fn scalar_value_string(value: Option<&AnyValue>) -> Option { match value.and_then(|v| v.value.as_ref())? { any_value::Value::StringValue(s) => Some(s.clone()), any_value::Value::IntValue(v) => Some(v.to_string()), any_value::Value::DoubleValue(v) => Some(v.to_string()), + any_value::Value::BoolValue(v) => Some(v.to_string()), _ => None, } } @@ -1535,6 +1536,15 @@ mod tests { } } + fn bool_keyvalue(key: &str, value: bool) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(Val::BoolValue(value)), + }), + } + } + fn descriptor_ctx() -> OtlpMetricCtx { OtlpMetricCtx { resource_info: true, @@ -1698,6 +1708,64 @@ mod tests { assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); } + #[test] + fn test_boolean_attributes_keep_series_distinct() { + set_default_prefix(None).unwrap(); + let request = metrics_request(vec![Metric { + name: "bool_repro".to_string(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![ + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", false)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(100.0)), + ..Default::default() + }, + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", true)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(7.0)), + ..Default::default() + }, + ], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + is_monotonic: true, + })), + ..Default::default() + }]); + let conversion = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + let rows = conversion + .requests + .inserts + .iter() + .find(|insert| insert.table_name == "bool_repro_total") + .expect("missing bool_repro_total table") + .rows + .as_ref() + .unwrap(); + assert_eq!(2, rows.rows.len()); + let failed = rows + .schema + .iter() + .position(|column| column.column_name == "failed") + .expect("boolean attribute `failed` was dropped from the label set"); + let values = rows + .rows + .iter() + .map(|row| row.values[failed].value_data.clone()) + .collect::>(); + assert_eq!( + vec![ + Some(ValueData::StringValue("false".to_string())), + Some(ValueData::StringValue("true".to_string())), + ], + values + ); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index c816583f135..d7b581ccd05 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -26,6 +26,7 @@ use datatypes::arrow::array::{Float64Array, Int64Array, StringArray, TimestampMi use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; +use query::datafusion::QUERY_PARALLELISM_HINT; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; use rstest::rstest; use rstest_reuse::apply; @@ -193,6 +194,139 @@ async fn promql_query_as_batches( } } +#[tokio::test(flavor = "multi_thread")] +async fn histogram_fraction_binary_join_preserves_partitioned_rows() { + common_telemetry::init_default_ut_logging(); + + let distributed = crate::tests::create_distributed_instance( + "histogram_fraction_binary_join_preserves_partitioned_rows", + ) + .await; + let instance = distributed.frontend(); + + for (table, create) in [ + ( + "pa_promql_m", + r#"CREATE TABLE pa_promql_m ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + );"#, + ), + ( + "lg_hist_m_2", + r#"CREATE TABLE phy_hist_m_2 ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + ) ENGINE=metric WITH ('physical_metric_table' = 'true'); + CREATE TABLE lg_hist_m_2 ( + ts TIMESTAMP TIME INDEX, + host STRING, + shard STRING, + val DOUBLE, + PRIMARY KEY (host, shard) + ) PARTITION ON COLUMNS (shard) ( + shard < 'b', + shard >= 'b' AND shard < 'c', + shard >= 'c' AND shard < 'd', + shard >= 'd' + ) ENGINE=metric WITH ('on_physical_table' = 'phy_hist_m_2');"#, + ), + ] { + execute_all( + &instance, + &format!( + "{create}\nINSERT INTO {table} (ts, host, shard, val) VALUES\n (0, 'h', 'a', 1),\n (0, 'h', 'b', 1),\n (0, 'h', 'c', 1),\n (0, 'h', 'd', 1);" + ), + QueryContext::arc(), + ) + .await; + + let expr = format!( + r#"histogram_fraction(-Inf,+Inf,abs(label_replace(({table} and on(host) (sum by(host)({table}))),"le","+Inf","missing_label",".*")))"# + ); + for parallelism in [1, 4] { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + let query_ctx = Arc::new(query_ctx); + for (query, expected_value) in [ + (format!("TQL EVAL (0, 0, '1s') {expr}"), 1.0), + ( + format!("TQL EVAL (0, 0, '1s') {expr} + on(host, shard) {table}"), + 2.0, + ), + ] { + let output = instance + .do_query(&query, query_ctx.clone()) + .await + .remove(0) + .unwrap(); + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + other => panic!("expected result batches, got {other:?}"), + }; + let mut rows = batches + .iter() + .flat_map(|batch| { + let ts = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let host_array = + cast(batch.column_by_name("host").unwrap(), &DataType::Utf8).unwrap(); + let host = host_array.as_any().downcast_ref::().unwrap(); + let shard_array = + cast(batch.column_by_name("shard").unwrap(), &DataType::Utf8).unwrap(); + let shard = shard_array.as_any().downcast_ref::().unwrap(); + let value = batch + .columns() + .iter() + .find_map(|column| column.as_any().downcast_ref::()) + .unwrap(); + (0..batch.num_rows()) + .map(move |row| { + ( + ts.value(row), + host.value(row).to_owned(), + shard.value(row).to_owned(), + value.value(row), + ) + }) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.2.cmp(&right.2)); + assert_eq!( + rows, + ["a", "b", "c", "d"] + .into_iter() + .map(|shard| (0, "h".to_owned(), shard.to_owned(), expected_value)) + .collect::>(), + "query_parallelism={parallelism}, query={query}" + ); + } + } + } +} + const ANON_PROMQL_RATIO_REPRO_DB: &str = "repro_db"; const ANON_PROMQL_RATIO_REPRO_CREATE: &str = r#" diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index fae78b94f85..5de18c54526 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -117,6 +117,101 @@ macro_rules! repartition_tests { }; } +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_physical_metric_without_logical_table() { + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&StorageType::File); + let home_dir = create_temp_dir("repartition_physical_metric_without_logical_table"); + let cluster = + GreptimeDbClusterBuilder::new("test_repartition_physical_metric_without_logical_table") + .await + .with_datanodes(3) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) + .build(true) + .await; + + let instance = cluster.fe_instance(); + let query_ctx = QueryContext::arc(); + let sql = r#" + CREATE TABLE `physical_metric_without_logical`( + `ts` TIMESTAMP TIME INDEX, + `val` DOUBLE, + `host` STRING PRIMARY KEY + ) PARTITION ON COLUMNS (`host`) ( + `host` < 'm', + `host` >= 'm' + ) ENGINE = metric + WITH ( + "physical_metric_table" = "true" + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_id = get_table_id(&cluster.metasrv, "physical_metric_without_logical").await; + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + let column_ids_before_split = table_info.table_info.meta.column_ids.clone(); + assert_eq!(column_ids_before_split, vec![0, 1, 2]); + + let sql = r#" + ALTER TABLE `physical_metric_without_logical` SPLIT PARTITION ( + `host` < 'm' + ) INTO ( + `host` < 'g', + `host` >= 'g' AND `host` < 'm' + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + assert_eq!( + table_info.table_info.meta.column_ids, + column_ids_before_split + ); + + let result = run_sql( + instance, + &query_partitions_sql("physical_metric_without_logical"), + query_ctx, + ) + .await + .unwrap(); + let expected = r#"+---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| table_catalog | table_schema | table_name | partition_name | partition_expression | partition_description | greptime_partition_id | partition_ordinal_position | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| greptime | public | physical_metric_without_logical | p0 | host | host < g | 4398046511104 | 1 | +| greptime | public | physical_metric_without_logical | p1 | host | host >= m | 4398046511105 | 2 | +| greptime | public | physical_metric_without_logical | p2 | host | host >= g AND host < m | 4398046511106 | 3 | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+"#; + check_output_stream(result.data, expected).await; +} + /// COUNT must use visible rows rather than the full row counts of shared SSTs. #[tokio::test(flavor = "multi_thread")] async fn test_repartition_append_count_file() {