diff --git a/config/config.md b/config/config.md
index 612a5cd43c0..4bce212b443 100644
--- a/config/config.md
+++ b/config/config.md
@@ -393,6 +393,11 @@
| `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". |
| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". |
| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". |
+| `write_retry` | -- | -- | Retry transient distributed write failures during region migration or repartition.
Backoff grows exponentially, with jitter to spread concurrent retries. |
+| `write_retry.max_retries` | Integer | `3` | Maximum retries after the initial write attempt. Set to 0 to disable retries. |
+| `write_retry.initial_backoff` | String | `1500ms` | Initial backoff before jitter is applied. Must be positive. |
+| `write_retry.max_backoff` | String | `10s` | Maximum backoff before jitter is applied. Must be at least initial_backoff. |
+| `write_retry.backoff_multiplier` | Float | `2.0` | Backoff multiplier. Must be finite and at least 1. |
| `datanode` | -- | -- | Datanode options. |
| `datanode.client` | -- | -- | Datanode client options. |
| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. |
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index e4ef6e36935..758625a5f1c 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -402,6 +402,18 @@ memory_pool_size = "50%"
#+ experimental_spill_compression = "uncompressed"
# --- End spill-to-disk ---
+## Retry transient distributed write failures during region migration or repartition.
+## Backoff grows exponentially, with jitter to spread concurrent retries.
+[write_retry]
+## Maximum retries after the initial write attempt. Set to 0 to disable retries.
+max_retries = 3
+## Initial backoff before jitter is applied. Must be positive.
+initial_backoff = "1500ms"
+## Maximum backoff before jitter is applied. Must be at least initial_backoff.
+max_backoff = "10s"
+## Backoff multiplier. Must be finite and at least 1.
+backoff_multiplier = 2.0
+
## Datanode options.
[datanode]
## Datanode client options.
diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs
index b000f7e1f02..8e07d128d48 100644
--- a/src/cmd/src/frontend.rs
+++ b/src/cmd/src/frontend.rs
@@ -512,10 +512,11 @@ impl StartCommand {
cached_meta_backend.clone(),
layered_cache_registry.clone(),
catalog_manager,
- client,
+ client.clone(),
meta_client.clone(),
process_manager,
- );
+ )
+ .with_write_retry_clients(client);
plugins::setup_frontend_plugins_post_build(&mut plugins, &plugin_opts, &builder)
.await
diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md
index fdaa5dbcfe4..8b6e22e02dd 100644
--- a/src/frontend/AGENTS.md
+++ b/src/frontend/AGENTS.md
@@ -48,7 +48,9 @@ including caches owned by pipeline and protocol handlers.
- **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` →
`check_permission` → `operator`'s `Inserter` (schema validation, optional
auto-create, partition routing, meter admission) → local `RegionServer`
- (standalone) or RPC to datanodes (distributed). Arrow bulk inserts pass the
+ (standalone) or RPC to datanodes (distributed). `FrontendBuilder` configures
+ the operator write dispatcher with frontend retry options for distributed
+ inserts, deletes, and bulk writes. Arrow bulk inserts pass the
request channel to `Inserter` and check meter admission for each nonempty batch.
- Finite ingestion requests split internally admit their total rows per database
before dispatch (`operator::insert::admit_write` / `admit_row_insert_batches`).
diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs
index 9fd2312e736..451040c884a 100644
--- a/src/frontend/src/error.rs
+++ b/src/frontend/src/error.rs
@@ -323,6 +323,13 @@ pub enum Error {
name: String,
},
+ #[snafu(display("Invalid write retry config: {}", reason))]
+ InvalidWriteRetryConfig {
+ reason: &'static str,
+ #[snafu(implicit)]
+ location: Location,
+ },
+
#[snafu(display("Invalid tls config"))]
InvalidTlsConfig {
#[snafu(source)]
@@ -391,7 +398,8 @@ impl ErrorExt for Error {
| Error::IllegalAuthConfig { .. }
| Error::ColumnNoneDefaultValue { .. }
| Error::IncompleteGrpcRequest { .. }
- | Error::InvalidTlsConfig { .. } => StatusCode::InvalidArguments,
+ | Error::InvalidTlsConfig { .. }
+ | Error::InvalidWriteRetryConfig { .. } => StatusCode::InvalidArguments,
Error::NotSupported { .. } => StatusCode::Unsupported,
diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs
index 90e7abc4517..5e6f0c13b22 100644
--- a/src/frontend/src/frontend.rs
+++ b/src/frontend/src/frontend.rs
@@ -36,9 +36,10 @@ use crate::heartbeat::HeartbeatTask;
use crate::instance::Instance;
use crate::service_config::{
InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions,
- PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions,
+ PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, WriteRetryOptions,
};
+/// Configures frontend protocols, query execution, and distributed write retries.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
pub struct FrontendOptions {
@@ -76,6 +77,8 @@ pub struct FrontendOptions {
pub meta_client: Option,
pub logging: LoggingOptions,
pub datanode: DatanodeClientOptions,
+ /// Retry policy for transient distributed write failures.
+ pub write_retry: WriteRetryOptions,
pub user_provider: Option,
pub tracing: TracingOptions,
pub query: QueryOptions,
@@ -113,6 +116,7 @@ impl Default for FrontendOptions {
meta_client: None,
logging: LoggingOptions::default(),
datanode: DatanodeClientOptions::default(),
+ write_retry: WriteRetryOptions::default(),
user_provider: None,
tracing: TracingOptions::default(),
query: QueryOptions::default(),
@@ -228,6 +232,72 @@ mod tests {
type GrpcStream =
Pin> + Send + Sync + 'static>>;
+ #[test]
+ fn test_write_retry_config() {
+ let defaults: FrontendOptions = toml::from_str("").unwrap();
+ assert_eq!(defaults.write_retry, WriteRetryOptions::default());
+ let options: FrontendOptions = toml::from_str(
+ r#"
+[write_retry]
+max_retries = 5
+initial_backoff = "250ms"
+max_backoff = "5s"
+backoff_multiplier = 3.0
+"#,
+ )
+ .unwrap();
+ assert_eq!(
+ options.write_retry,
+ WriteRetryOptions {
+ max_retries: 5,
+ initial_backoff: Duration::from_millis(250),
+ max_backoff: Duration::from_secs(5),
+ backoff_multiplier: 3.0,
+ }
+ );
+ options.write_retry.validate().unwrap();
+ let serialized = toml::to_string(&options).unwrap();
+ let parsed: FrontendOptions = toml::from_str(&serialized).unwrap();
+ assert_eq!(options.write_retry, parsed.write_retry);
+ let disabled: FrontendOptions = toml::from_str("[write_retry]\nmax_retries = 0").unwrap();
+ assert_eq!(disabled.write_retry.max_retries, 0);
+ assert_eq!(
+ disabled.write_retry.initial_backoff,
+ defaults.write_retry.initial_backoff
+ );
+ }
+
+ #[test]
+ fn test_write_retry_config_from_env() {
+ temp_env::with_vars(
+ [
+ ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_RETRIES", Some("4")),
+ (
+ "FRONTEND_RETRY_TEST__WRITE_RETRY__INITIAL_BACKOFF",
+ Some("500ms"),
+ ),
+ ("FRONTEND_RETRY_TEST__WRITE_RETRY__MAX_BACKOFF", Some("8s")),
+ (
+ "FRONTEND_RETRY_TEST__WRITE_RETRY__BACKOFF_MULTIPLIER",
+ Some("1.5"),
+ ),
+ ],
+ || {
+ let options =
+ FrontendOptions::load_layered_options(None, "FRONTEND_RETRY_TEST").unwrap();
+ assert_eq!(
+ options.write_retry,
+ WriteRetryOptions {
+ max_retries: 4,
+ initial_backoff: Duration::from_millis(500),
+ max_backoff: Duration::from_secs(8),
+ backoff_multiplier: 1.5,
+ }
+ );
+ },
+ );
+ }
+
#[test]
fn test_logical_batcher_protocols_from_env() {
temp_env::with_vars(
diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs
index 032722499f9..0d35eff6044 100644
--- a/src/frontend/src/instance/builder.rs
+++ b/src/frontend/src/instance/builder.rs
@@ -21,11 +21,14 @@ use catalog::CatalogManagerRef;
use catalog::kvbackend::KvBackendCatalogManager;
use catalog::process_manager::ProcessManagerRef;
use catalog::system_schema::semantic_graph::EntityGraphProviderRef;
+use client::client_manager::NodeClients;
use common_base::Plugins;
use common_batcher::request_limiter::RequestLimiter;
use common_datasource::object_store::LocalFileAccess;
use common_event_recorder::{EventRecorderImpl, EventRecorderRef};
-use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef};
+use common_meta::cache::{
+ LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableInfoCacheRef, TableRouteCacheRef,
+};
use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator};
use common_meta::key::TableMetadataManager;
use common_meta::key::flow::FlowMetadataManager;
@@ -46,6 +49,7 @@ use operator::statement::{
StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef,
};
use operator::table::TableMutationOperator;
+use operator::write_dispatcher::WriteDispatcher;
use partition::cache::PartitionInfoCacheRef;
use partition::manager::PartitionRuleManager;
use pipeline::pipeline_operator::PipelineOperator;
@@ -73,6 +77,7 @@ pub struct FrontendBuilder {
local_cache_invalidator: Option,
catalog_manager: CatalogManagerRef,
node_manager: NodeManagerRef,
+ write_retry_clients: Option>,
plugins: Option,
procedure_executor: ProcedureExecutorRef,
process_manager: ProcessManagerRef,
@@ -80,6 +85,7 @@ pub struct FrontendBuilder {
}
impl FrontendBuilder {
+ /// Creates a frontend builder from its configuration and shared services.
#[allow(clippy::too_many_arguments)]
pub fn new(
options: FrontendOptions,
@@ -97,6 +103,7 @@ impl FrontendBuilder {
local_cache_invalidator: None,
catalog_manager,
node_manager,
+ write_retry_clients: None,
plugins: None,
procedure_executor,
process_manager,
@@ -168,6 +175,14 @@ impl FrontendBuilder {
&self.catalog_manager
}
+ /// Enables write retries through the distributed datanode clients.
+ pub fn with_write_retry_clients(self, clients: Arc) -> Self {
+ Self {
+ write_retry_clients: Some(clients),
+ ..self
+ }
+ }
+
pub fn node_manager(&self) -> &NodeManagerRef {
&self.node_manager
}
@@ -187,7 +202,11 @@ impl FrontendBuilder {
}
}
+ /// Validates configuration and assembles the frontend [`Instance`].
pub async fn try_build(self) -> Result {
+ if let Err(reason) = self.options.write_retry.validate() {
+ return error::InvalidWriteRetryConfigSnafu { reason }.fail();
+ }
let kv_backend = self.kv_backend;
let node_manager = self.node_manager;
let plugins = self.plugins.unwrap_or_default();
@@ -214,6 +233,24 @@ impl FrontendBuilder {
.local_cache_invalidator
.unwrap_or_else(|| Arc::new(DummyCacheInvalidator));
+ let mut write_dispatcher =
+ WriteDispatcher::new(node_manager.clone(), partition_manager.clone());
+ if let Some(clients) = self.write_retry_clients {
+ let table_info_cache: TableInfoCacheRef =
+ self.layered_cache_registry
+ .get()
+ .context(error::CacheRequiredSnafu {
+ name: cache::TABLE_INFO_CACHE_NAME,
+ })?;
+ write_dispatcher = write_dispatcher.with_retry(
+ clients,
+ self.options.write_retry.clone(),
+ table_info_cache,
+ local_cache_invalidator.clone(),
+ );
+ }
+ let write_dispatcher = Arc::new(write_dispatcher);
+
let region_query_handler =
if let Some(factory) = plugins.get::() {
factory.build(partition_manager.clone(), node_manager.clone())
@@ -236,6 +273,7 @@ impl FrontendBuilder {
table_flownode_cache.clone(),
self.options.auto_create_table,
)
+ .with_write_dispatcher(write_dispatcher.clone())
};
// The execution-only inserter owns no batchers, avoiding an Arc cycle.
let bulk_inserter = Arc::new(create_inserter());
@@ -264,11 +302,14 @@ impl FrontendBuilder {
create_inserter()
.with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())),
);
- let deleter = Arc::new(Deleter::new(
- self.catalog_manager.clone(),
- partition_manager.clone(),
- node_manager.clone(),
- ));
+ let deleter = Arc::new(
+ Deleter::new(
+ self.catalog_manager.clone(),
+ partition_manager.clone(),
+ node_manager.clone(),
+ )
+ .with_write_dispatcher(write_dispatcher),
+ );
let requester = Arc::new(Requester::new(
self.catalog_manager.clone(),
partition_manager.clone(),
diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs
index 255766ac14d..423f2996cae 100644
--- a/src/frontend/src/instance/logical_batcher.rs
+++ b/src/frontend/src/instance/logical_batcher.rs
@@ -48,6 +48,7 @@ impl Instance {
LogicalTablePendingRowsBatcher::try_new(
self.partition_manager().clone(),
self.node_manager().clone(),
+ self.inserter.write_dispatcher().clone(),
self.catalog_manager().clone(),
self.table_flownode_set_cache().clone(),
true,
diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs
index 2c761c4242a..535ae9b655b 100644
--- a/src/frontend/src/service_config.rs
+++ b/src/frontend/src/service_config.rs
@@ -25,6 +25,7 @@ pub use influxdb::{InfluxdbMergeMode, InfluxdbOptions};
pub use jaeger::JaegerOptions;
pub use mysql::MysqlOptions;
pub use opentsdb::OpentsdbOptions;
+pub use operator::write_dispatcher::WriteRetryOptions;
pub use otlp::OtlpOptions;
pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions};
pub use postgres::PostgresOptions;
diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs
index dd20f34adf8..b7ad7646e3c 100644
--- a/src/metric-engine/src/engine/bulk_insert.rs
+++ b/src/metric-engine/src/engine/bulk_insert.rs
@@ -354,17 +354,32 @@ mod tests {
#[tokio::test]
async fn test_bulk_insert_physical_region_passthrough() {
- check_bulk_insert_physical_region_passthrough(false, false).await;
- check_bulk_insert_physical_region_passthrough(true, false).await;
+ check_bulk_insert_physical_region_passthrough(false, false, false).await;
+ check_bulk_insert_physical_region_passthrough(true, false, false).await;
}
#[tokio::test]
async fn test_bulk_insert_physical_region_close_recovery() {
- check_bulk_insert_physical_region_passthrough(false, true).await;
- check_bulk_insert_physical_region_passthrough(true, true).await;
+ check_bulk_insert_physical_region_passthrough(false, true, false).await;
+ check_bulk_insert_physical_region_passthrough(true, true, false).await;
}
- async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) {
+ #[tokio::test]
+ async fn test_physical_bulk_partition_columns_preserve_stored_rows() {
+ for skip_wal in [false, true] {
+ let stripped =
+ check_bulk_insert_physical_region_passthrough(skip_wal, true, false).await;
+ let retained =
+ check_bulk_insert_physical_region_passthrough(skip_wal, true, true).await;
+ assert_eq!(retained, stripped, "skip_wal={skip_wal}");
+ }
+ }
+
+ async fn check_bulk_insert_physical_region_passthrough(
+ skip_wal: bool,
+ close: bool,
+ retain_partition_columns: bool,
+ ) -> String {
// Use flat format so that BulkMemtable is used (supports write_bulk).
let mito_config = MitoConfig {
default_flat_format: true,
@@ -394,7 +409,11 @@ mod tests {
index: 2,
column_id: 2, // column_id for "job" in the physical table
}];
- let non_tag_indices = vec![0, 1]; // timestamp, value
+ let non_tag_indices = if retain_partition_columns {
+ vec![0, 1, 2] // timestamp, value, and the job partition column
+ } else {
+ vec![0, 1] // timestamp, value
+ };
let second_batch = build_logical_batch(3, 3);
let physical_batch = modify_batch_sparse(
second_batch,
@@ -403,6 +422,10 @@ mod tests {
&non_tag_indices,
)
.unwrap();
+ assert_eq!(
+ physical_batch.column_by_name("job").is_some(),
+ retain_partition_columns
+ );
let request = build_bulk_request(physical_region_id, physical_batch, skip_wal);
let response = env
.metric()
@@ -469,6 +492,7 @@ mod tests {
batches.iter().map(|b| b.num_rows()).sum::(),
if skip_wal && !close { 0 } else { 6 },
);
+ batches.pretty_print().unwrap()
}
#[tokio::test]
diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md
index d20704cf616..26751a7b02e 100644
--- a/src/servers/AGENTS.md
+++ b/src/servers/AGENTS.md
@@ -43,3 +43,6 @@ cargo nextest run -p servers
Keep protocol translation here and permissions/database behavior in frontend.
Async table batching returns storage-completion futures with its early
acknowledgements. Keep completion reporting independent of the caller's lifetime.
+Logical-table physical flushes retain partition columns in Arrow RPC payloads and
+pass all region writes to the operator write dispatcher together; the dispatcher
+owns peer resolution and retries. See `batcher/logical_table/region_write.rs`.
diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs
index be992692d13..d879486c2d3 100644
--- a/src/servers/benches/flush_batch_physical.rs
+++ b/src/servers/benches/flush_batch_physical.rs
@@ -15,8 +15,6 @@
use std::collections::HashMap;
use std::sync::Arc;
-use api::region::RegionResponse;
-use api::v1::meta::Peer;
use api::v1::region::RegionRequest;
use arrow::array::{Float64Array, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit};
@@ -29,11 +27,10 @@ use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
use partition::error::Result as PartitionResult;
use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
use servers::batcher::logical_table::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical,
};
-use servers::error::{self, Result};
-use store_api::storage::RegionId;
+use servers::error;
use table::test_util::table_info::test_table_info;
use tokio::runtime::Runtime;
@@ -95,24 +92,17 @@ impl PhysicalFlushPartitionProvider for BenchPartitionProvider {
async fn find_table_partition_rule(
&self,
_table_info: &table::metadata::TableInfo,
- ) -> PartitionResult {
- Ok(Arc::new(SingleRegionPartitionRule))
- }
-
- async fn find_region_leader(&self, _region_id: RegionId) -> Result {
- Ok(Peer {
- id: 1,
- addr: "bench-node".to_string(),
- })
+ ) -> PartitionResult<(PartitionRuleRef, HashMap>)> {
+ Ok((Arc::new(SingleRegionPartitionRule), HashMap::new()))
}
}
-struct BenchNodeRequester;
+struct BenchWriteDispatcher;
#[async_trait]
-impl PhysicalFlushNodeRequester for BenchNodeRequester {
- async fn handle(&self, _peer: &Peer, _request: RegionRequest) -> error::Result {
- Ok(RegionResponse::new(0))
+impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
+ Ok(0)
}
}
@@ -224,7 +214,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
table: make_physical_table_metadata(num_tags),
};
let partition_provider = BenchPartitionProvider;
- let node_requester = BenchNodeRequester;
+ let write_dispatcher = BenchWriteDispatcher;
let mut group = c.benchmark_group("flush_batch_physical");
@@ -244,7 +234,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
"phy",
&ctx,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
)
.await
@@ -271,7 +261,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) {
"phy",
&ctx,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
)
.await
diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs
index 7e0cea44b6a..f34b0ae406e 100644
--- a/src/servers/src/batcher/logical_table.rs
+++ b/src/servers/src/batcher/logical_table.rs
@@ -39,6 +39,7 @@ use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
use common_time::timestamp::TimeUnit;
use meter_core::data::MeterRecord;
use meter_macros::write_meter;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use snafu::ResultExt;
@@ -52,7 +53,7 @@ use crate::batcher::logical_table::pending_worker::{
PendingWorker, WorkerCommand, remove_worker_if_same_channel, start_worker,
};
pub use crate::batcher::logical_table::region_write::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata,
};
pub use crate::batcher::logical_table::tables::{
@@ -97,7 +98,7 @@ pub struct LogicalTablePendingRowsBatcher {
flush_interval: Duration,
flush_policy: TimingFlushPolicy,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
flush_limiter: FlushLimiter,
@@ -110,10 +111,12 @@ pub struct LogicalTablePendingRowsBatcher {
}
impl LogicalTablePendingRowsBatcher {
+ /// Creates the batcher when metric-engine batching is enabled.
#[allow(clippy::too_many_arguments)]
pub fn try_new(
partition_manager: PartitionRuleManagerRef,
node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
table_flownode_set_cache: TableFlownodeSetCacheRef,
prom_store_with_metric_engine: bool,
@@ -153,7 +156,7 @@ impl LogicalTablePendingRowsBatcher {
flush_interval,
flush_policy,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
prom_store_with_metric_engine,
@@ -404,7 +407,7 @@ impl LogicalTablePendingRowsBatcher {
rx,
self.shutdown.clone(),
self.partition_manager.clone(),
- self.node_manager.clone(),
+ self.write_dispatcher.clone(),
self.catalog_manager.clone(),
self.flow_notification_tx.clone(),
worker_idle_timeout,
diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs
index ac2db137de1..1ce90f39d01 100644
--- a/src/servers/src/batcher/logical_table/batch.rs
+++ b/src/servers/src/batcher/logical_table/batch.rs
@@ -18,45 +18,41 @@ use std::time::Instant;
use catalog::CatalogManagerRef;
use common_batcher::flush_limiter::FlushLimiter;
-use common_meta::node_manager::NodeManagerRef;
-use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
+use common_query::prelude;
use common_telemetry::{debug, warn};
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use snafu::OptionExt;
use crate::batcher::flow_notifier::FlowNotifier;
-use crate::batcher::logical_table::PHYSICAL_TABLE_KEY;
-use crate::batcher::logical_table::batch_convert::{
- TableBatch, concat_modified_batches, transform_logical_batches_to_physical,
-};
-use crate::batcher::logical_table::flow_notifier::enqueue_flow_notifications;
+use crate::batcher::logical_table;
+use crate::batcher::logical_table::batch_convert::{self, TableBatch};
+use crate::batcher::logical_table::flow_notifier;
use crate::batcher::logical_table::pending_worker::FlushWaiter;
use crate::batcher::logical_table::region_write::{
- CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter,
- PartitionManagerPhysicalFlushAdapter, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester,
- PhysicalFlushPartitionProvider, encode_region_write_requests, flush_region_writes_concurrently,
- plan_region_batches, resolve_region_targets,
+ self, CatalogManagerPhysicalFlushAdapter, PartitionManagerPhysicalFlushAdapter,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
+ WriteDispatcherPhysicalFlushAdapter,
};
use crate::error;
use crate::error::Result;
-use crate::metrics::{
- FLUSH_DROPPED_ROWS, FLUSH_ELAPSED, FLUSH_FAILURES, FLUSH_ROWS, FLUSH_TOTAL,
- PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED,
-};
+use crate::metrics;
-pub(in crate::batcher::logical_table) struct Batch {
- pub(in crate::batcher::logical_table) table_batches: Vec,
- pub(in crate::batcher::logical_table) total_row_count: usize,
- pub(in crate::batcher::logical_table) db_string: String,
- pub(in crate::batcher::logical_table) ctx: QueryContextRef,
- pub(in crate::batcher::logical_table) waiters: Vec,
+/// Rows and completion waiters owned by one physical flush.
+pub(crate) struct Batch {
+ pub(crate) table_batches: Vec,
+ pub(crate) total_row_count: usize,
+ pub(crate) db_string: String,
+ pub(crate) ctx: QueryContextRef,
+ pub(crate) waiters: Vec,
}
-pub(in crate::batcher::logical_table) async fn spawn_flush(
+/// Spawns a physical flush with the configured concurrency limit.
+pub(crate) async fn spawn_flush(
flush: Batch,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
flush_limiter: FlushLimiter,
@@ -68,7 +64,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
flush_batch_with_managers(
flush,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
)
@@ -80,7 +76,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
flush_batch_with_managers(
flush,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
)
@@ -89,32 +85,32 @@ pub(in crate::batcher::logical_table) async fn spawn_flush(
}
}
-pub(in crate::batcher::logical_table) async fn flush_batch_with_managers(
+/// Connects shared routing and dispatch dependencies to a physical flush.
+pub(crate) async fn flush_batch_with_managers(
flush: Batch,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
) {
let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager };
- let node_requester = NodeManagerPhysicalFlushAdapter {
- node_manager: node_manager.clone(),
- };
+ let write_dispatcher = WriteDispatcherPhysicalFlushAdapter { write_dispatcher };
let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager };
flush_batch(
flush,
&partition_provider,
- &node_requester,
+ &write_dispatcher,
&catalog_provider,
flow_notification_tx,
)
.await;
}
-pub(in crate::batcher::logical_table) async fn flush_batch(
+/// Flushes a batch and reports completion to its waiters.
+pub(crate) async fn flush_batch(
flush: Batch,
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
+ write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
flow_notification_tx: FlowNotifier,
) {
@@ -130,21 +126,21 @@ pub(in crate::batcher::logical_table) async fn flush_batch(
// Physical-table-level flush: transform all logical table batches
// into physical format and write them together.
let physical_table_name = ctx
- .extension(PHYSICAL_TABLE_KEY)
- .unwrap_or(GREPTIME_PHYSICAL_TABLE)
+ .extension(logical_table::PHYSICAL_TABLE_KEY)
+ .unwrap_or(prelude::GREPTIME_PHYSICAL_TABLE)
.to_string();
let result = flush_batch_physical(
&table_batches,
&physical_table_name,
&ctx,
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
)
.await;
let elapsed = start.elapsed().as_secs_f64();
- FLUSH_ELAPSED.observe(elapsed);
+ metrics::FLUSH_ELAPSED.observe(elapsed);
debug!(
"Pending rows batch flushed, total rows: {}, elapsed time: {}s",
@@ -153,18 +149,18 @@ pub(in crate::batcher::logical_table) async fn flush_batch(
match result {
Ok(affected_rows) => {
- FLUSH_TOTAL.inc();
- FLUSH_ROWS.observe(total_row_count as f64);
+ metrics::FLUSH_TOTAL.inc();
+ metrics::FLUSH_ROWS.observe(total_row_count as f64);
operator::metrics::DIST_INGEST_ROW_COUNT
.with_label_values(&[db_string.as_str()])
.inc_by(affected_rows as u64);
notify_waiters(waiters, Ok(()));
- enqueue_flow_notifications(table_batches, &flow_notification_tx);
+ flow_notifier::enqueue_flow_notifications(table_batches, &flow_notification_tx);
}
Err(err) => {
- FLUSH_FAILURES.inc();
- FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64);
+ metrics::FLUSH_FAILURES.inc();
+ metrics::FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64);
notify_waiters(waiters, Err(err));
}
}
@@ -185,12 +181,12 @@ pub async fn flush_batch_physical(
physical_table_name: &str,
ctx: &QueryContextRef,
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
+ write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
) -> Result {
// 1. Resolve the physical table and get column ID mapping
let physical_table = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_resolve_table"])
.start_timer();
catalog_manager
@@ -220,8 +216,8 @@ pub async fn flush_batch_physical(
})?;
// 2. Get the physical table's partition rule (one lookup instead of N)
- let partition_rule = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let (partition_rule, partition_versions) = {
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_fetch_partition_rule"])
.start_timer();
partition_manager
@@ -233,24 +229,29 @@ pub async fn flush_batch_physical(
partition_columns.iter().map(String::as_str).collect();
// 3. Transform each logical table batch into physical format
- let modified_batches =
- transform_logical_batches_to_physical(table_batches, &name_to_ids, &partition_columns_set)?;
+ let modified_batches = batch_convert::transform_logical_batches_to_physical(
+ table_batches,
+ &name_to_ids,
+ &partition_columns_set,
+ )?;
// 4. Concatenate all modified batches (all share the same physical schema)
- let combined_batch = concat_modified_batches(&modified_batches)?;
+ let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?;
// 5. Split by physical partition rule and send to regions
let physical_table_id = physical_table_info.table_id();
- let planned_batches = plan_region_batches(
+ let planned_batches = region_write::plan_region_batches(
combined_batch,
physical_table_id,
partition_rule.as_ref(),
- partition_columns,
)?;
- let resolved_batches = resolve_region_targets(planned_batches, partition_manager).await?;
- let region_writes = encode_region_write_requests(resolved_batches, ctx.skip_wal())?;
- flush_region_writes_concurrently(node_manager, region_writes).await
+ let region_writes = region_write::encode_region_write_requests(
+ planned_batches,
+ ctx.skip_wal(),
+ &partition_versions,
+ )?;
+ write_dispatcher.dispatch(region_writes).await
}
pub(in crate::batcher::logical_table) fn notify_waiters(
@@ -276,7 +277,6 @@ mod tests {
use std::task::Poll;
use std::time::Duration;
- use api::region::RegionResponse;
use api::v1::meta::Peer;
use api::v1::region::RegionRequest;
use arrow::record_batch::RecordBatch;
@@ -288,7 +288,6 @@ mod tests {
use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
use partition::error::Result as PartitionResult;
use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
- use store_api::storage::RegionId;
use table::metadata::TableId;
use table::test_util::table_info::test_table_info;
use tokio::sync::{mpsc, oneshot};
@@ -300,7 +299,7 @@ mod tests {
use crate::batcher::logical_table::batch_convert::TableBatch;
use crate::batcher::logical_table::pending_worker::FlushWaiter;
use crate::batcher::logical_table::region_write::{
- PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+ PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
PhysicalTableMetadata,
};
use crate::batcher::logical_table::test_util::{
@@ -344,9 +343,8 @@ mod tests {
},
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &MockFlushNodeRequester {
+ &MockFlushWriteDispatcher {
writes: writes.clone(),
fail: false,
},
@@ -405,9 +403,8 @@ mod tests {
},
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &MockFlushNodeRequester {
+ &MockFlushWriteDispatcher {
writes: writes.clone(),
fail: true,
},
@@ -451,8 +448,7 @@ mod tests {
row_count: 1,
}];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
flush_batch_physical(
@@ -461,9 +457,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -472,8 +467,7 @@ mod tests {
.unwrap();
assert_eq!(1, partition_calls.load(Ordering::SeqCst));
- assert_eq!(1, leader_calls.load(Ordering::SeqCst));
- assert_eq!(1, node.writes.load(Ordering::SeqCst));
+ assert_eq!(1, dispatcher.writes.load(Ordering::SeqCst));
}
#[tokio::test]
@@ -492,9 +486,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: Arc::new(AtomicUsize::new(0)),
- region_leader_calls: Arc::new(AtomicUsize::new(0)),
},
- &AffectedRowsFlushNodeRequester { affected_rows: 7 },
+ &AffectedRowsFlushWriteDispatcher { affected_rows: 7 },
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -506,7 +499,7 @@ mod tests {
}
#[tokio::test]
- async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() {
+ async fn test_flush_batch_physical_stops_before_partition_and_dispatch_when_table_missing() {
let table_batches = vec![TableBatch {
table_name: "t1".to_string(),
table_id: 11,
@@ -514,8 +507,7 @@ mod tests {
row_count: 1,
}];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
let err = flush_batch_physical(
@@ -524,9 +516,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider { table: None },
)
.await
@@ -537,8 +528,7 @@ mod tests {
.contains("Physical table 'missing_phy' not found")
);
assert_eq!(0, partition_calls.load(Ordering::SeqCst));
- assert_eq!(0, leader_calls.load(Ordering::SeqCst));
- assert_eq!(0, node.writes.load(Ordering::SeqCst));
+ assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
}
#[tokio::test]
@@ -558,8 +548,7 @@ mod tests {
},
];
let partition_calls = Arc::new(AtomicUsize::new(0));
- let leader_calls = Arc::new(AtomicUsize::new(0));
- let node = MockFlushNodeRequester::default();
+ let dispatcher = MockFlushWriteDispatcher::default();
let ctx = session::context::QueryContext::arc();
let err = flush_batch_physical(
@@ -568,9 +557,8 @@ mod tests {
&ctx,
&MockFlushPartitionProvider {
partition_rule_calls: partition_calls.clone(),
- region_leader_calls: leader_calls.clone(),
},
- &node,
+ &dispatcher,
&MockFlushCatalogProvider {
table: Some(mock_physical_table_metadata(1024)),
},
@@ -580,8 +568,7 @@ mod tests {
assert!(err.to_string().contains("unknown_tag"));
assert_eq!(1, partition_calls.load(Ordering::SeqCst));
- assert_eq!(0, leader_calls.load(Ordering::SeqCst));
- assert_eq!(0, node.writes.load(Ordering::SeqCst));
+ assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
}
fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata {
@@ -671,7 +658,6 @@ mod tests {
struct MockFlushPartitionProvider {
partition_rule_calls: Arc,
- region_leader_calls: Arc,
}
#[async_trait]
@@ -679,40 +665,28 @@ mod tests {
async fn find_table_partition_rule(
&self,
_table_info: &table::metadata::TableInfo,
- ) -> PartitionResult {
+ ) -> PartitionResult<(PartitionRuleRef, HashMap>)> {
self.partition_rule_calls.fetch_add(1, Ordering::SeqCst);
- Ok(Arc::new(SingleRegionPartitionRule))
- }
-
- async fn find_region_leader(&self, _region_id: RegionId) -> error::Result {
- self.region_leader_calls.fetch_add(1, Ordering::SeqCst);
- Ok(Peer {
- id: 1,
- addr: "node-1".to_string(),
- })
+ Ok((Arc::new(SingleRegionPartitionRule), HashMap::new()))
}
}
#[derive(Default)]
- struct MockFlushNodeRequester {
+ struct MockFlushWriteDispatcher {
writes: Arc,
fail: bool,
}
#[async_trait]
- impl PhysicalFlushNodeRequester for MockFlushNodeRequester {
- async fn handle(
- &self,
- _peer: &Peer,
- _request: RegionRequest,
- ) -> error::Result {
+ impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
self.writes.fetch_add(1, Ordering::SeqCst);
if self.fail {
return Err(Error::Internal {
err_msg: "physical write failed".to_string(),
});
}
- Ok(RegionResponse::new(0))
+ Ok(0)
}
}
@@ -726,18 +700,14 @@ mod tests {
}
#[derive(Default)]
- struct AffectedRowsFlushNodeRequester {
+ struct AffectedRowsFlushWriteDispatcher {
affected_rows: usize,
}
#[async_trait]
- impl PhysicalFlushNodeRequester for AffectedRowsFlushNodeRequester {
- async fn handle(
- &self,
- _peer: &Peer,
- _request: RegionRequest,
- ) -> error::Result {
- Ok(RegionResponse::new(self.affected_rows))
+ impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher {
+ async fn dispatch(&self, _requests: Vec) -> error::Result {
+ Ok(self.affected_rows)
}
}
}
diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs
index 37080dbe258..65204ec038c 100644
--- a/src/servers/src/batcher/logical_table/batch_convert.rs
+++ b/src/servers/src/batcher/logical_table/batch_convert.rs
@@ -29,8 +29,6 @@ use crate::error;
use crate::error::Result;
use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
-pub(in crate::batcher::logical_table) const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3;
-
/// An aligned logical record batch and its timestamp column index.
#[derive(Debug, Clone)]
pub struct RecordBatchWithTsIdx {
@@ -188,23 +186,6 @@ pub(in crate::batcher::logical_table) fn columns_taxonomy(
Ok((tag_columns, essential_column_indices))
}
-pub(in crate::batcher::logical_table) fn strip_partition_columns_from_batch(
- batch: RecordBatch,
-) -> Result {
- ensure!(
- batch.num_columns() >= PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
- error::InternalSnafu {
- err_msg: format!(
- "Expected at least {} columns in physical batch, got {}",
- PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
- batch.num_columns()
- ),
- }
- );
- let essential_indices: Vec = (0..PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT).collect();
- batch.project(&essential_indices).context(error::ArrowSnafu)
-}
-
/// Transforms logical table batches into physical format (sparse primary key encoding).
///
/// It identifies tag columns and essential columns (timestamp, value) for each logical batch
@@ -290,79 +271,18 @@ mod tests {
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
- use arrow::array::{Array, BinaryArray, StringArray, TimestampMillisecondArray};
+ use arrow::array::{Array, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
use metric_engine::batch_modifier::modify_batch_sparse;
use smallvec::SmallVec;
use crate::batcher::logical_table::batch_convert::{
- TableBatch, columns_taxonomy, strip_partition_columns_from_batch,
- transform_logical_batches_to_physical,
+ TableBatch, columns_taxonomy, transform_logical_batches_to_physical,
};
use crate::batcher::logical_table::test_util::mock_aligned_tag_batch;
use crate::error::Error;
- #[test]
- fn test_strip_partition_columns_from_batch_removes_partition_tags() {
- let batch = RecordBatch::try_new(
- Arc::new(ArrowSchema::new(vec![
- Field::new("__primary_key", ArrowDataType::Binary, false),
- Field::new(
- "greptime_timestamp",
- ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("greptime_value", ArrowDataType::Float64, true),
- Field::new("host", ArrowDataType::Utf8, true),
- ])),
- vec![
- Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
- Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
- Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
- Arc::new(StringArray::from(vec!["node-1"])),
- ],
- )
- .unwrap();
-
- let stripped = strip_partition_columns_from_batch(batch).unwrap();
-
- assert_eq!(3, stripped.num_columns());
- assert_eq!("__primary_key", stripped.schema().field(0).name());
- assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
- assert_eq!("greptime_value", stripped.schema().field(2).name());
- }
-
- #[test]
- fn test_strip_partition_columns_from_batch_projects_essential_columns_without_lookup() {
- let batch = RecordBatch::try_new(
- Arc::new(ArrowSchema::new(vec![
- Field::new("__primary_key", ArrowDataType::Binary, false),
- Field::new(
- "greptime_timestamp",
- ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("greptime_value", ArrowDataType::Float64, true),
- Field::new("host", ArrowDataType::Utf8, true),
- ])),
- vec![
- Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
- Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
- Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
- Arc::new(StringArray::from(vec!["node-1"])),
- ],
- )
- .unwrap();
-
- let stripped = strip_partition_columns_from_batch(batch).unwrap();
-
- assert_eq!(3, stripped.num_columns());
- assert_eq!("__primary_key", stripped.schema().field(0).name());
- assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
- assert_eq!("greptime_value", stripped.schema().field(2).name());
- }
-
#[test]
fn test_collect_tag_columns_and_non_tag_indices_keeps_partition_tag_column() {
let schema = Arc::new(ArrowSchema::new(vec![
diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs
index 4a83d8a87f4..6f46c717588 100644
--- a/src/servers/src/batcher/logical_table/pending_worker.rs
+++ b/src/servers/src/batcher/logical_table/pending_worker.rs
@@ -22,8 +22,8 @@ use common_batcher::flush_policy::FlushTrigger;
use common_batcher::flush_policy::timing::TimingFlushPolicy;
use common_batcher::pending_worker::PendingWorker as PendingCore;
use common_batcher::worker_registry::WorkerRegistry;
-use common_meta::node_manager::NodeManagerRef;
use common_telemetry::debug;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use table::metadata::TableId;
@@ -31,7 +31,7 @@ use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot};
use crate::batcher::flow_notifier::FlowNotifier;
use crate::batcher::logical_table::BatchKey;
-use crate::batcher::logical_table::batch::{Batch, flush_batch_with_managers, spawn_flush};
+use crate::batcher::logical_table::batch::{self, Batch};
use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch};
use crate::error::Error;
use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS};
@@ -53,7 +53,8 @@ pub(in crate::batcher::logical_table) struct PendingWorker {
pub(in crate::batcher::logical_table) tx: mpsc::Sender,
}
-pub(in crate::batcher::logical_table) enum WorkerCommand {
+/// Commands accepted by a logical-table batching worker.
+pub(crate) enum WorkerCommand {
Submit {
table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>,
total_rows: usize,
@@ -92,15 +93,16 @@ impl PendingBatch {
}
}
+/// Starts a worker that accumulates logical rows and dispatches physical flushes.
#[allow(clippy::too_many_arguments)]
-pub(in crate::batcher::logical_table) fn start_worker(
+pub(crate) fn start_worker(
key: BatchKey,
worker_tx: mpsc::Sender,
workers: Arc>,
mut rx: mpsc::Receiver,
shutdown: broadcast::Sender<()>,
partition_manager: PartitionRuleManagerRef,
- node_manager: NodeManagerRef,
+ write_dispatcher: WriteDispatcherRef,
catalog_manager: CatalogManagerRef,
flow_notification_tx: FlowNotifier,
worker_idle_timeout: Duration,
@@ -140,10 +142,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
PENDING_ROWS.add(total_rows as i64);
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) {
- spawn_flush(
+ batch::spawn_flush(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
flush_limiter.clone(),
@@ -152,10 +154,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
None => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
- flush_batch_with_managers(
+ batch::flush_batch_with_managers(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
).await;
@@ -189,10 +191,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
_ = pending_flush.wait_flush() => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) {
- spawn_flush(
+ batch::spawn_flush(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
flush_limiter.clone(),
@@ -201,10 +203,10 @@ pub(in crate::batcher::logical_table) fn start_worker(
}
_ = shutdown_rx.recv() => {
if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
- flush_batch_with_managers(
+ batch::flush_batch_with_managers(
flush,
partition_manager.clone(),
- node_manager.clone(),
+ write_dispatcher.clone(),
catalog_manager.clone(),
flow_notification_tx.clone(),
).await;
@@ -502,6 +504,10 @@ mod tests {
let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager {
datanodes: Arc::new(HashMap::new()),
});
+ let write_dispatcher = Arc::new(operator::write_dispatcher::WriteDispatcher::new(
+ node_manager,
+ partition_manager.clone(),
+ ));
let catalog_manager = MemoryCatalogManager::with_default_setup();
let (flow_notification_tx, _flow_notification_rx) =
FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
@@ -515,7 +521,7 @@ mod tests {
worker_rx,
shutdown.clone(),
partition_manager,
- node_manager,
+ write_dispatcher,
catalog_manager,
flow_notification_tx,
worker_idle_timeout,
diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs
index 71b558d1428..34bc2b7bf46 100644
--- a/src/servers/src/batcher/logical_table/region_write.rs
+++ b/src/servers/src/batcher/logical_table/region_write.rs
@@ -14,29 +14,27 @@
use std::collections::HashMap;
-use api::v1::ArrowIpc;
-use api::v1::meta::Peer;
use api::v1::region::{
BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
};
-use arrow::compute::filter_record_batch;
+use api::v1::{ArrowIpc, PartitionExprVersion};
+use arrow::compute;
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use catalog::CatalogManagerRef;
use common_grpc::error::Error as GrpcError;
-use common_grpc::flight::record_batch_to_ipc;
-use common_meta::node_manager::NodeManagerRef;
+use common_grpc::flight;
use common_telemetry::tracing_context::TracingContext;
+use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use partition::partition::PartitionRuleRef;
use snafu::ResultExt;
use store_api::storage::RegionId;
use table::metadata::{TableId, TableInfo, TableInfoRef};
-use crate::batcher::logical_table::batch_convert::strip_partition_columns_from_batch;
use crate::error;
use crate::error::{Error, Result};
-use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
+use crate::metrics;
#[derive(Clone)]
pub struct PhysicalTableMetadata {
@@ -56,23 +54,21 @@ pub trait PhysicalFlushCatalogProvider: Send + Sync {
) -> catalog::error::Result>;
}
+/// Loads the physical partition rule and its expression versions.
#[async_trait]
pub trait PhysicalFlushPartitionProvider: Send + Sync {
+ /// Loads a routing snapshot for the physical table.
async fn find_table_partition_rule(
&self,
table_info: &TableInfo,
- ) -> partition::error::Result;
-
- async fn find_region_leader(&self, region_id: RegionId) -> Result;
+ ) -> partition::error::Result<(PartitionRuleRef, HashMap>)>;
}
+/// Dispatches the region requests produced by a physical flush.
#[async_trait]
-pub trait PhysicalFlushNodeRequester: Send + Sync {
- async fn handle(
- &self,
- peer: &Peer,
- request: RegionRequest,
- ) -> Result;
+pub trait PhysicalFlushWriteDispatcher: Send + Sync {
+ /// Dispatches one physical flush together, sharing routing refreshes across its writes.
+ async fn dispatch(&self, requests: Vec) -> Result;
}
#[derive(Clone)]
@@ -115,125 +111,61 @@ impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter {
async fn find_table_partition_rule(
&self,
table_info: &TableInfo,
- ) -> partition::error::Result {
+ ) -> partition::error::Result<(PartitionRuleRef, HashMap>)> {
self.partition_manager
.find_table_partition_rule(table_info)
.await
- .map(|(rule, _)| rule)
- }
-
- async fn find_region_leader(&self, region_id: RegionId) -> Result {
- let peer = self.partition_manager.find_region_leader(region_id).await?;
- Ok(peer)
}
}
+/// Sends a physical flush through the shared write dispatcher.
#[derive(Clone)]
-pub(in crate::batcher::logical_table) struct NodeManagerPhysicalFlushAdapter {
- pub(in crate::batcher::logical_table) node_manager: NodeManagerRef,
+pub(crate) struct WriteDispatcherPhysicalFlushAdapter {
+ pub(crate) write_dispatcher: WriteDispatcherRef,
}
#[async_trait]
-impl PhysicalFlushNodeRequester for NodeManagerPhysicalFlushAdapter {
- async fn handle(
- &self,
- peer: &Peer,
- request: RegionRequest,
- ) -> error::Result {
- let datanode = self.node_manager.datanode(peer).await;
- datanode
- .handle(request)
+impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter {
+ async fn dispatch(&self, requests: Vec) -> Result {
+ self.write_dispatcher
+ .dispatch_with_metrics(
+ requests,
+ metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ .with_label_values(&["flush_write_region"]),
+ )
.await
.context(error::CommonMetaSnafu)
+ .map(|response| response.affected_rows)
}
}
-pub(in crate::batcher::logical_table) struct FlushRegionWrite {
- pub(in crate::batcher::logical_table) datanode: Peer,
- pub(in crate::batcher::logical_table) request: RegionRequest,
-}
-
-pub(in crate::batcher::logical_table) struct PlannedRegionBatch {
- pub(in crate::batcher::logical_table) region_id: RegionId,
- pub(in crate::batcher::logical_table) batch: RecordBatch,
+/// Arrow rows selected for one physical region.
+pub(crate) struct PlannedRegionBatch {
+ pub(crate) region_id: RegionId,
+ pub(crate) batch: RecordBatch,
}
#[cfg(test)]
impl PlannedRegionBatch {
- pub(in crate::batcher::logical_table) fn num_rows(&self) -> usize {
+ pub(crate) fn num_rows(&self) -> usize {
self.batch.num_rows()
}
}
-pub(in crate::batcher::logical_table) struct ResolvedRegionBatch {
- pub(in crate::batcher::logical_table) planned: PlannedRegionBatch,
- pub(in crate::batcher::logical_table) datanode: Peer,
-}
-
-pub(in crate::batcher::logical_table) fn should_dispatch_concurrently(
- region_write_count: usize,
-) -> bool {
- region_write_count > 1
-}
-
-pub(in crate::batcher::logical_table) async fn flush_region_writes_concurrently(
- node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
- writes: Vec,
-) -> Result {
- let mut affected_rows = 0;
- if !should_dispatch_concurrently(writes.len()) {
- for write in writes {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_write_region"])
- .start_timer();
- affected_rows += node_manager
- .handle(&write.datanode, write.request)
- .await?
- .affected_rows;
- }
- return Ok(affected_rows);
- }
-
- let write_futures = writes.into_iter().map(|write| async move {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_write_region"])
- .start_timer();
-
- let response = node_manager.handle(&write.datanode, write.request).await?;
- Ok::<_, Error>(response.affected_rows)
- });
-
- // todo(hl): should be bounded.
- let affected_rows = futures::future::try_join_all(write_futures)
- .await?
- .into_iter()
- .sum();
- Ok(affected_rows)
-}
-
pub(in crate::batcher::logical_table) fn split_combined_batch_by_region(
combined_batch: &RecordBatch,
partition_rule: &dyn partition::partition::PartitionRule,
) -> Result> {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_split_record_batch"])
.start_timer();
let map = partition_rule.split_record_batch(combined_batch)?;
Ok(map)
}
-pub(in crate::batcher::logical_table) fn prepare_physical_region_routing_batch(
- combined_batch: RecordBatch,
- partition_columns: &[String],
-) -> Result {
- if partition_columns.is_empty() {
- return Ok(combined_batch);
- }
- strip_partition_columns_from_batch(combined_batch)
-}
-
-pub(in crate::batcher::logical_table) fn plan_region_batch(
- stripped_batch: &RecordBatch,
+/// Selects the rows assigned to one physical region.
+pub(crate) fn plan_region_batch(
+ combined_batch: &RecordBatch,
physical_table_id: TableId,
region_number: u32,
mask: &partition::partition::RegionMask,
@@ -243,12 +175,12 @@ pub(in crate::batcher::logical_table) fn plan_region_batch(
}
let region_batch = if mask.select_all() {
- stripped_batch.clone()
+ combined_batch.clone()
} else {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_filter_record_batch"])
.start_timer();
- filter_record_batch(stripped_batch, mask.array()).context(error::ArrowSnafu)?
+ compute::filter_record_batch(combined_batch, mask.array()).context(error::ArrowSnafu)?
};
let row_count = region_batch.num_rows();
@@ -262,19 +194,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batch(
}))
}
-pub(in crate::batcher::logical_table) fn plan_region_batches(
+/// Partitions a physical flush while retaining columns needed for rerouting.
+pub(crate) fn plan_region_batches(
combined_batch: RecordBatch,
physical_table_id: TableId,
partition_rule: &dyn partition::partition::PartitionRule,
- partition_columns: &[String],
) -> Result> {
let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?;
- let stripped_batch = prepare_physical_region_routing_batch(combined_batch, partition_columns)?;
+ // Retain partition columns in the RPC payload so routing can be rebuilt after repartition.
+ // Mito's sparse bulk conversion projects the stored fields and encoded primary key.
let mut planned_batches = Vec::new();
for (region_number, mask) in region_masks {
if let Some(planned_batch) =
- plan_region_batch(&stripped_batch, physical_table_id, region_number, &mask)?
+ plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)?
{
planned_batches.push(planned_batch);
}
@@ -283,39 +216,20 @@ pub(in crate::batcher::logical_table) fn plan_region_batches(
Ok(planned_batches)
}
-pub(in crate::batcher::logical_table) async fn resolve_region_targets(
+/// Encodes each region batch with the partition version used to select its rows.
+pub(crate) fn encode_region_write_requests(
planned_batches: Vec,
- partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
-) -> Result> {
- let mut resolved_batches = Vec::with_capacity(planned_batches.len());
- for planned in planned_batches {
- let datanode = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
- .with_label_values(&["flush_physical_resolve_region_leader"])
- .start_timer();
- partition_manager
- .find_region_leader(planned.region_id)
- .await?
- };
-
- resolved_batches.push(ResolvedRegionBatch { planned, datanode });
- }
-
- Ok(resolved_batches)
-}
-
-pub(in crate::batcher::logical_table) fn encode_region_write_requests(
- resolved_batches: Vec,
skip_wal: bool,
-) -> Result> {
- let mut region_writes = Vec::with_capacity(resolved_batches.len());
- for resolved in resolved_batches {
- let region_id = resolved.planned.region_id;
+ partition_versions: &HashMap>,
+) -> Result> {
+ let mut region_writes = Vec::with_capacity(planned_batches.len());
+ for planned in planned_batches {
+ let region_id = planned.region_id;
let (schema_bytes, data_header, payload) = {
- let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
+ let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
.with_label_values(&["flush_physical_encode_ipc"])
.start_timer();
- record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)?
+ flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)?
};
let request = RegionRequest {
@@ -326,7 +240,11 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests(
body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
skip_wal,
region_id: region_id.as_u64(),
- partition_expr_version: None,
+ partition_expr_version: partition_versions
+ .get(®ion_id.region_number())
+ .copied()
+ .flatten()
+ .map(|value| PartitionExprVersion { value }),
// Set aligned_schema_version to None so that datanode will check the batch schema again to see if any
// column is missing.
aligned_schema_version: None,
@@ -338,10 +256,7 @@ pub(in crate::batcher::logical_table) fn encode_region_write_requests(
})),
};
- region_writes.push(FlushRegionWrite {
- datanode: resolved.datanode,
- request,
- });
+ region_writes.push(request);
}
Ok(region_writes)
@@ -359,86 +274,21 @@ pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Erro
#[cfg(test)]
mod tests {
- use std::collections::HashMap;
- use std::sync::Arc;
- use std::sync::atomic::{AtomicUsize, Ordering};
- use std::time::Duration;
-
- use api::region::RegionResponse;
- use api::v1::meta::Peer;
- use api::v1::region::{RegionRequest, region_request};
- use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
+ use api::v1::region::region_request;
+ use arrow::array::{Array, BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
- use async_trait::async_trait;
- use common_meta::error::Result as MetaResult;
- use common_meta::node_manager::{Datanode, DatanodeRef};
- use common_query::request::QueryRequest;
- use common_recordbatch::SendableRecordBatchStream;
use partition::partition::{PartitionRule, RegionMask};
+ use std::collections::HashMap;
+ use std::sync::Arc;
use store_api::storage::RegionId;
- use tokio::time::sleep;
use crate::batcher::logical_table::region_write::{
- FlushRegionWrite, PlannedRegionBatch, ResolvedRegionBatch, encode_region_write_requests,
- flush_region_writes_concurrently, plan_region_batches, should_dispatch_concurrently,
+ PlannedRegionBatch, encode_region_write_requests, plan_region_batches,
};
- use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager;
-
- #[tokio::test]
- async fn test_flush_region_writes_concurrently_dispatches_multiple_datanodes() {
- let inflight = Arc::new(AtomicUsize::new(0));
- let max_inflight = Arc::new(AtomicUsize::new(0));
- let datanode1: DatanodeRef = Arc::new(ConcurrentMockDatanode {
- delay: Duration::from_millis(100),
- inflight: inflight.clone(),
- max_inflight: max_inflight.clone(),
- });
- let datanode2: DatanodeRef = Arc::new(ConcurrentMockDatanode {
- delay: Duration::from_millis(100),
- inflight,
- max_inflight: max_inflight.clone(),
- });
-
- let mut datanodes = HashMap::new();
- datanodes.insert(1, datanode1);
- datanodes.insert(2, datanode2);
- let node_manager = Arc::new(ConcurrentMockNodeManager {
- datanodes: Arc::new(datanodes),
- });
-
- let writes = vec![
- FlushRegionWrite {
- datanode: Peer {
- id: 1,
- addr: "node1".to_string(),
- },
- request: RegionRequest::default(),
- },
- FlushRegionWrite {
- datanode: Peer {
- id: 2,
- addr: "node2".to_string(),
- },
- request: RegionRequest::default(),
- },
- ];
-
- flush_region_writes_concurrently(node_manager.as_ref(), writes)
- .await
- .unwrap();
- assert!(max_inflight.load(Ordering::SeqCst) >= 2);
- }
#[test]
- fn test_should_dispatch_concurrently_by_region_count() {
- assert!(!should_dispatch_concurrently(0));
- assert!(!should_dispatch_concurrently(1));
- assert!(should_dispatch_concurrently(2));
- }
-
- #[test]
- fn test_plan_region_batches_splits_and_strips_partition_columns() {
+ fn test_plan_region_batches_preserves_partition_columns_through_encoding() {
let combined_batch = RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
Field::new("__primary_key", ArrowDataType::Binary, false),
@@ -464,7 +314,6 @@ mod tests {
&TwoRegionPartitionRule {
partition_columns: vec!["host".to_string()],
},
- &["host".to_string()],
)
.unwrap();
planned_batches.sort_by_key(|planned| planned.region_id.region_number());
@@ -472,10 +321,42 @@ mod tests {
assert_eq!(2, planned_batches.len());
assert_eq!(RegionId::new(1024, 1), planned_batches[0].region_id);
assert_eq!(1, planned_batches[0].num_rows());
- assert_eq!(3, planned_batches[0].batch.num_columns());
+ assert_eq!(4, planned_batches[0].batch.num_columns());
assert_eq!(RegionId::new(1024, 2), planned_batches[1].region_id);
assert_eq!(1, planned_batches[1].num_rows());
- assert_eq!(3, planned_batches[1].batch.num_columns());
+ assert_eq!(4, planned_batches[1].batch.num_columns());
+
+ let writes = encode_region_write_requests(
+ planned_batches,
+ false,
+ &HashMap::from([(1, Some(7)), (2, Some(8))]),
+ )
+ .unwrap();
+ for (index, write) in writes.into_iter().enumerate() {
+ let Some(region_request::Body::BulkInsert(request)) = write.body else {
+ panic!("expected bulk")
+ };
+ assert_eq!(
+ request.partition_expr_version.unwrap().value,
+ 7 + index as u64
+ );
+ let Some(api::v1::region::bulk_insert_request::Body::ArrowIpc(ipc)) = request.body
+ else {
+ panic!("expected IPC")
+ };
+ let batch = common_grpc::flight::FlightDecoder::try_from_schema_bytes(&ipc.schema)
+ .unwrap()
+ .try_decode_record_batch(&ipc.data_header, &ipc.payload)
+ .unwrap();
+ let tags = batch
+ .column_by_name("host")
+ .unwrap()
+ .as_any()
+ .downcast_ref::()
+ .unwrap();
+ assert_eq!(tags.len(), 1);
+ assert_eq!(tags.value(0), if index == 0 { "node-1" } else { "node-2" });
+ }
}
#[test]
@@ -520,44 +401,6 @@ mod tests {
}
}
- #[derive(Clone)]
- struct ConcurrentMockDatanode {
- delay: Duration,
- inflight: Arc,
- max_inflight: Arc,
- }
-
- #[async_trait]
- impl Datanode for ConcurrentMockDatanode {
- async fn handle(&self, _request: RegionRequest) -> MetaResult {
- let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1;
- loop {
- let max = self.max_inflight.load(Ordering::SeqCst);
- if now <= max {
- break;
- }
- if self
- .max_inflight
- .compare_exchange(max, now, Ordering::SeqCst, Ordering::SeqCst)
- .is_ok()
- {
- break;
- }
- }
-
- sleep(self.delay).await;
- self.inflight.fetch_sub(1, Ordering::SeqCst);
- Ok(RegionResponse::new(0))
- }
-
- async fn handle_query(
- &self,
- _request: QueryRequest,
- ) -> MetaResult {
- unimplemented!()
- }
- }
-
fn check_encode_region_write_requests(skip_wal: bool) {
let planned_batch = PlannedRegionBatch {
region_id: RegionId::new(1024, 1),
@@ -579,21 +422,22 @@ mod tests {
)
.unwrap(),
};
- let resolved_batch = ResolvedRegionBatch {
- planned: planned_batch,
- datanode: Peer {
- id: 1,
- addr: "node-1".to_string(),
- },
- };
- let writes = encode_region_write_requests(vec![resolved_batch], skip_wal).unwrap();
+ let writes = encode_region_write_requests(
+ vec![planned_batch],
+ skip_wal,
+ &HashMap::from([(1, Some(7))]),
+ )
+ .unwrap();
assert_eq!(1, writes.len());
- assert_eq!(1, writes[0].datanode.id);
- let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.body else {
+ let Some(region_request::Body::BulkInsert(request)) = &writes[0].body else {
panic!("expected bulk insert request");
};
assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id);
assert_eq!(skip_wal, request.skip_wal);
+ assert_eq!(
+ Some(api::v1::PartitionExprVersion { value: 7 }),
+ request.partition_expr_version
+ );
}
}
diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs
index 1cc81f9b615..491bd6ff2b3 100644
--- a/src/servers/src/batcher/logical_table/test_util.rs
+++ b/src/servers/src/batcher/logical_table/test_util.rs
@@ -15,10 +15,9 @@
use std::collections::HashMap;
use std::sync::Arc;
-use api::region::RegionResponse;
use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse};
use api::v1::meta::Peer;
-use api::v1::region::{InsertRequests, RegionRequest};
+use api::v1::region::InsertRequests;
use arrow::array::{StringArray, TimestampMillisecondArray};
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
@@ -27,12 +26,9 @@ use common_meta::error::Result as MetaResult;
use common_meta::node_manager::{
DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef,
};
-use snafu::ResultExt;
use tokio::sync::mpsc;
use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx;
-use crate::batcher::logical_table::region_write::PhysicalFlushNodeRequester;
-use crate::error;
pub(in crate::batcher::logical_table) fn mock_tag_batch(
tag_name: &str,
@@ -149,14 +145,3 @@ impl FlownodeManager for FlowNotificationMockNodeManager {
self.flownode.clone()
}
}
-
-#[async_trait]
-impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager {
- async fn handle(&self, peer: &Peer, request: RegionRequest) -> error::Result {
- let datanode = self.datanode(peer).await;
- datanode
- .handle(request)
- .await
- .context(error::CommonMetaSnafu)
- }
-}