refactor: relocate the logical table batcher

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/frontend/src/instance/prom_store.rs

@@ -42,11 +42,11 @@ use operator::insert::{
 use operator::statement::StatementExecutor;
 use prost::Message;
 use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
+use servers::batcher::logical_table::PendingRowsSchemaAlterer;
 use servers::error::{self, AuthSnafu, Result as ServerResult};
 use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics};
 use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
 use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
-use servers::pending_rows_batcher::PendingRowsSchemaAlterer;
 use servers::prom_store;
 use servers::query_handler::{
     PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse,

src/frontend/src/server.rs

@@ -24,6 +24,9 @@ use common_base::Plugins;
 use common_config::Configurable;
 use common_telemetry::{info, warn};
 use meta_client::MetaClientOptions;
+use servers::batcher::logical_table::{
+    LogicalTablePendingRowsBatcher, pending_rows_batch_sync_enabled,
+};
 use servers::error::Error as ServerError;
 use servers::grpc::builder::GrpcServerBuilder;
 use servers::grpc::flight::FlightCraftRef;
@@ -38,7 +41,6 @@ use servers::interceptor::LogIngestInterceptorRef;
 use servers::metrics_handler::MetricsHandler;
 use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
 use servers::otel_arrow::OtelArrowServiceHandler;
-use servers::pending_rows_batcher::{PendingRowsBatcher, pending_rows_batch_sync_enabled};
 use servers::postgres::PostgresServer;
 use servers::request_memory_limiter::ServerMemoryLimiter;
 use servers::server::{Server, ServerHandlers};
@@ -132,7 +134,7 @@ where
         let prom_store = effective_prom_store_options(opts);
         if prom_store.enable {
             let pending_rows_batcher = if prom_store.with_metric_engine {
-                PendingRowsBatcher::try_new(
+                LogicalTablePendingRowsBatcher::try_new(
                     self.instance.partition_manager().clone(),
                     self.instance.node_manager().clone(),
                     self.instance.catalog_manager().clone(),
@@ -715,7 +717,7 @@ mod tests {
 
     #[test]
     fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() {
-        // Mirrors the conditions under which `PendingRowsBatcher::try_new`
+        // Mirrors the conditions under which `LogicalTablePendingRowsBatcher::try_new`
         // returns `None`; in these cases no request can wait for a pending-row
         // flush, so the timeout must not be raised.
         type KnobMutator = fn(&mut FrontendOptions);

src/servers/benches/flush_batch_physical.rs

@@ -28,11 +28,11 @@ use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
 use partition::error::Result as PartitionResult;
 use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
-use servers::error::{self, Result};
-use servers::pending_rows_batcher::{
+use servers::batcher::logical_table::{
     PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
     PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical,
 };
+use servers::error::{self, Result};
 use store_api::storage::RegionId;
 use table::test_util::table_info::test_table_info;
 use tokio::runtime::Runtime;

src/servers/src/batcher.rs

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Ordinary-table batching and shared notification components.
+//! Logical-table and ordinary-table batching implementations.
 
 mod flow_notifier;
 mod flow_sender;
+pub mod logical_table;
 pub mod table;
 
 #[cfg(test)]

src/servers/src/batcher/logical_table.rs

@@ -346,7 +346,7 @@ fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey {
 }
 
 /// Prometheus remote write pending rows batcher.
-pub struct PendingRowsBatcher {
+pub struct LogicalTablePendingRowsBatcher {
     workers: Arc<WorkerRegistry<BatchKey, WorkerCommand>>,
     flush_interval: Duration,
     flush_policy: TimingFlushPolicy,
@@ -363,7 +363,7 @@ pub struct PendingRowsBatcher {
     shutdown: broadcast::Sender<()>,
 }
 
-impl PendingRowsBatcher {
+impl LogicalTablePendingRowsBatcher {
     #[allow(clippy::too_many_arguments)]
     pub fn try_new(
         partition_manager: PartitionRuleManagerRef,
@@ -832,6 +832,7 @@ impl PendingRowsBatcher {
 
         Ok(aligned_batches)
     }
+
     async fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker {
         let (tx, receiver) = self
             .workers
@@ -872,7 +873,7 @@ impl PendingRowsBatcher {
     }
 }
 
-impl Drop for PendingRowsBatcher {
+impl Drop for LogicalTablePendingRowsBatcher {
     fn drop(&mut self) {
         let _ = self.shutdown.send(());
     }
@@ -1887,22 +1888,22 @@ mod tests {
     use tokio::sync::{Notify, Semaphore, broadcast, mpsc, oneshot};
     use tokio::time::{advance, sleep};
 
-    use crate::error;
-    use crate::metrics::FLOW_NOTIFICATION_DROPPED;
-    use crate::pending_rows_batcher::{
+    use crate::batcher::logical_table::{
         BatchKey, Error, FlushBatch, FlushLimiter, FlushRegionWrite, FlushTrigger, FlushWaiter,
-        Notifier, PendingBatch, PendingCore, PendingRowsBatcher, PhysicalFlushCatalogProvider,
-        PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata,
-        PlannedRegionBatch, RecordBatchWithTsIdx, RequestLimiter, ResolvedRegionBatch, TableBatch,
-        TimingFlushPolicy, WorkerCommand, WorkerRegistry, batch_key_from_ctx, columns_taxonomy,
-        drain_batch, encode_region_write_requests, extract_timestamps, flush_batch,
-        flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp,
+        LogicalTablePendingRowsBatcher, Notifier, PendingBatch, PendingCore,
+        PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider,
+        PhysicalTableMetadata, PlannedRegionBatch, RecordBatchWithTsIdx, RequestLimiter,
+        ResolvedRegionBatch, TableBatch, TimingFlushPolicy, WorkerCommand, WorkerRegistry,
+        batch_key_from_ctx, columns_taxonomy, drain_batch, encode_region_write_requests, extract_timestamps,
+        flush_batch, flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp,
         notify_flow_dirty_windows_after_flush, notify_waiters, plan_region_batches,
         remove_worker_if_same_channel, should_close_worker_on_idle_timeout,
         should_dispatch_concurrently, start_flow_notification_worker, start_worker,
         strip_partition_columns_from_batch, transform_logical_batches_to_physical,
         try_enqueue_flow_notification,
     };
+    use crate::error;
+    use crate::metrics::FLOW_NOTIFICATION_DROPPED;
     use crate::prom_row_builder::rows_to_aligned_record_batch;
 
     fn mock_rows(row_count: usize, schema_name: &str) -> Rows {
@@ -2127,7 +2128,7 @@ mod tests {
     #[test]
     fn test_flow_notification_queue_drops_when_full() {
         let (tx, mut rx) = Notifier::try_new(1).unwrap();
-        let notification = |table_id| crate::pending_rows_batcher::FlowNotification {
+        let notification = |table_id| crate::batcher::logical_table::FlowNotification {
             table_id,
             timestamps: vec![table_id as i64],
         };
@@ -2334,7 +2335,8 @@ mod tests {
             ],
         };
 
-        let (table_rows, total_rows) = PendingRowsBatcher::collect_non_empty_table_rows(requests);
+        let (table_rows, total_rows) =
+            LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests);
 
         assert_eq!(2, total_rows);
         assert_eq!(1, table_rows.len());
@@ -2601,7 +2603,7 @@ mod tests {
     fn mock_flow_notification_sender(
         cache: TableFlownodeSetCacheRef,
         node_manager: NodeManagerRef,
-    ) -> Notifier<crate::pending_rows_batcher::FlowNotification> {
+    ) -> Notifier<crate::batcher::logical_table::FlowNotification> {
         let (tx, rx) = Notifier::try_new(16).unwrap();
         start_flow_notification_worker(rx, cache, node_manager);
         tx

src/servers/src/http.rs

@@ -53,6 +53,7 @@ use tower_http::trace::TraceLayer;
 
 use self::authorize::AuthState;
 use self::result::table_result::TableResponse;
+use crate::batcher::logical_table::LogicalTablePendingRowsBatcher;
 use crate::elasticsearch;
 use crate::error::{
     AddressBindSnafu, AlreadyStartedSnafu, Error, InternalIoSnafu, InvalidHeaderValueSnafu, Result,
@@ -74,7 +75,6 @@ use crate::http::result::null_result::NullResponse;
 use crate::interceptor::LogIngestInterceptorRef;
 use crate::metrics::http_metrics_layer;
 use crate::metrics_handler::MetricsHandler;
-use crate::pending_rows_batcher::PendingRowsBatcher;
 use crate::prometheus_handler::PrometheusHandlerRef;
 use crate::query_handler::sql::ServerSqlQueryHandlerRef;
 use crate::query_handler::{
@@ -727,7 +727,7 @@ impl HttpServerBuilder {
         prom_store_with_metric_engine: bool,
         prom_validation_mode: PromValidationMode,
         experimental_enable_prometheus_native_histogram: bool,
-        pending_rows_batcher: Option<Arc<PendingRowsBatcher>>,
+        pending_rows_batcher: Option<Arc<LogicalTablePendingRowsBatcher>>,
     ) -> Self {
         let state = PromStoreState {
             prom_store_handler: handler,

src/servers/src/http/prom_store.rs

@@ -40,12 +40,12 @@ use table::requests::{
     SOURCE_PROMETHEUS,
 };
 
+use crate::batcher::logical_table::LogicalTablePendingRowsBatcher;
 use crate::error::{self, InternalSnafu, PipelineSnafu, Result};
 use crate::http::extractor::PipelineInfo;
 use crate::http::header::{
     CONTENT_TYPE_PROTOBUF_STR, GREPTIME_DB_HEADER_METRICS, write_cost_header_map,
 };
-use crate::pending_rows_batcher::PendingRowsBatcher;
 use crate::prom_remote_write::decode::PromSeriesProcessor;
 use crate::prom_remote_write::v2::decode_remote_write_v2;
 use crate::prom_remote_write::validation::PromValidationMode;
@@ -75,7 +75,7 @@ pub struct PromStoreState {
     pub prom_store_with_metric_engine: bool,
     pub prom_validation_mode: PromValidationMode,
     pub experimental_enable_prometheus_native_histogram: bool,
-    pub pending_rows_batcher: Option<Arc<PendingRowsBatcher>>,
+    pub pending_rows_batcher: Option<Arc<LogicalTablePendingRowsBatcher>>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -358,9 +358,9 @@ trait PromWriteBatcher: Send + Sync {
 }
 
 #[async_trait]
-impl PromWriteBatcher for PendingRowsBatcher {
+impl PromWriteBatcher for LogicalTablePendingRowsBatcher {
     async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result<u64> {
-        PendingRowsBatcher::submit(self, requests, ctx).await
+        LogicalTablePendingRowsBatcher::submit(self, requests, ctx).await
     }
 }
 
@@ -387,7 +387,7 @@ async fn preflight_prometheus_rows(
 /// sample/histogram headers even when a later table write fails.
 async fn write_prometheus_rows_with_progress(
     prom_store_handler: PromStoreProtocolHandlerRef,
-    pending_rows_batcher: Option<Arc<PendingRowsBatcher>>,
+    pending_rows_batcher: Option<Arc<LogicalTablePendingRowsBatcher>>,
     prom_store_with_metric_engine: bool,
     mut batches: Vec<PromWriteBatch>,
 ) -> std::result::Result<PromWriteOutcome, PromWriteError> {
@@ -453,7 +453,7 @@ async fn write_prometheus_rows_with_progress(
 
 async fn write_prometheus_v2_rows_with_progress(
     prom_store_handler: PromStoreProtocolHandlerRef,
-    pending_rows_batcher: Option<Arc<PendingRowsBatcher>>,
+    pending_rows_batcher: Option<Arc<LogicalTablePendingRowsBatcher>>,
     prom_store_with_metric_engine: bool,
     sample_batches: Vec<PromWriteBatch>,
     histogram_batches: Vec<PromWriteBatch>,

src/servers/src/lib.rs

@@ -38,7 +38,6 @@ pub mod mysql;
 pub mod opentsdb;
 pub mod otel_arrow;
 pub mod otlp;
-pub mod pending_rows_batcher;
 mod pipeline;
 pub mod postgres;
 pub mod prom_remote_write;

src/servers/src/prom_row_builder.rs

@@ -35,9 +35,9 @@ use datatypes::data_type::DataType;
 use datatypes::prelude::ConcreteDataType;
 use snafu::{OptionExt, ResultExt, ensure};
 
+use crate::batcher::logical_table::RecordBatchWithTsIdx;
 use crate::error;
 use crate::error::Result;
-use crate::pending_rows_batcher::RecordBatchWithTsIdx;
 
 /// Extract timestamp, field, and tag column names from a logical region schema.
 fn unzip_logical_region_schema(

tests-integration/src/test_util.rs

@@ -58,6 +58,7 @@ use object_store::config::{
 use object_store::services::{Azblob, Gcs, Oss, S3};
 use object_store::test_util::TempFolder;
 use object_store::{AzblobConnection, GcsConnection, ObjectStore, OssConnection, S3Connection};
+use servers::batcher::logical_table::LogicalTablePendingRowsBatcher;
 use servers::grpc::builder::GrpcServerBuilder;
 use servers::grpc::greptime_handler::GreptimeRequestHandler;
 use servers::grpc::{FlightCompression, GrpcOptions, GrpcServer, GrpcServerConfig};
@@ -65,7 +66,6 @@ use servers::http::{HttpOptions, HttpServerBuilder};
 use servers::metrics_handler::MetricsHandler;
 use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
 use servers::otel_arrow::OtelArrowServiceHandler;
-use servers::pending_rows_batcher::PendingRowsBatcher;
 use servers::postgres::PostgresServer;
 use servers::prom_remote_write::validation::PromValidationMode;
 use servers::query_handler::sql::SqlQueryHandler;
@@ -867,7 +867,7 @@ pub fn build_test_prom_server(
     // Mirror the production wiring at `frontend::server`: build the batcher from the
     // instance's managers. A short flush interval keeps the test responsive.
     let pending_rows_batcher = if enable_batcher {
-        PendingRowsBatcher::try_new(
+        LogicalTablePendingRowsBatcher::try_new(
             frontend_ref.partition_manager().clone(),
             frontend_ref.node_manager().clone(),
             frontend_ref.catalog_manager().clone(),