diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 3990c72ebc3..1096ff262e1 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -409,9 +409,6 @@ impl Inserter { future::join_all(submissions).await }; let affected_rows = results.into_iter().sum::>()?; - crate::metrics::DIST_INGEST_ROW_COUNT - .with_label_values(&[ctx.get_db_string().as_str()]) - .inc_by(affected_rows as u64); Ok(Output::new( OutputData::AffectedRows(affected_rows), OutputMeta::new_with_cost(write_cost as _), diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 3b88a4ca827..1a46a7f10f3 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -225,6 +225,7 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -238,19 +239,23 @@ mod tests { use arrow::array::{Int32Array, TimestampMillisecondArray}; use arrow::datatypes::Schema as ArrowSchema; use arrow::record_batch::RecordBatch; + use catalog::memory::MemoryCatalogManager; use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; + use common_catalog::consts::default_engine; use common_grpc::flight::FlightDecoder; use common_meta::error::Result as MetaResult; use common_meta::peer::Peer; use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; + use common_query::OutputData; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use common_telemetry::info; use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder}; use datatypes::value::Value as DtValue; + use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef}; use operator::batcher::PendingRowsBatcher; use operator::error::Error; use operator::insert::Inserter; @@ -261,7 +266,9 @@ mod tests { }; use session::context::{Channel, QueryContext}; use store_api::storage::RegionId; + use table::dist_table::DistTable; use table::metadata::TableInfoRef; + use table::requests::InsertRequest; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::time::timeout; @@ -364,14 +371,15 @@ mod tests { expected_schema: schema_name.clone(), })); let inserter = Arc::new(Inserter::new( - catalog::memory::MemoryCatalogManager::new(), - partitions, - nodes, + MemoryCatalogManager::new(), + partitions.clone(), + nodes.clone(), mock_table_flownode_cache(1, vec![]).await, true, )); let mut table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter()); table.schema_name = schema_name; + table.meta.engine = default_engine().to_string(); let mut columns = table.meta.schema.column_schemas().to_vec(); columns[2] = columns[2] .clone() @@ -399,32 +407,67 @@ mod tests { inserter, ) .unwrap(); - let influx_ctx = Arc::new(QueryContext::with_channel( - "greptime", - &table.schema_name, - Channel::Influx, - )); - let opentsdb_ctx = Arc::new(QueryContext::with_channel( - "greptime", - &table.schema_name, - Channel::Opentsdb, - )); + let context = |channel| { + let mut ctx = QueryContext::with_channel("greptime", &table.schema_name, channel); + ctx.set_batching_enabled(true); + Arc::new(ctx) + }; + let influx_ctx = context(Channel::Influx); + let opentsdb_ctx = context(Channel::Opentsdb); let ingest_count = DIST_INGEST_ROW_COUNT.with_label_values(&[influx_ctx.get_db_string().as_str()]); assert_eq!(0, ingest_count.get()); influx_ctx.set_skip_wal(skip_wal); opentsdb_ctx.set_skip_wal(skip_wal); - let first_permit = batcher.acquire().await.unwrap(); - let second_permit = if shared_request { - first_permit.clone() - } else { - batcher.acquire().await.unwrap() - }; let (first_result, second_result) = timeout(Duration::from_secs(5), async { - tokio::join!( - batcher.submit(table.clone(), first, influx_ctx, first_permit), - batcher.submit(table.clone(), second, opentsdb_ctx, second_permit), - ) + if shared_request { + let permit = batcher.acquire().await.unwrap(); + tokio::join!( + batcher.submit(table.clone(), first, influx_ctx, permit.clone()), + batcher.submit(table.clone(), second, opentsdb_ctx, permit), + ) + } else { + let inserter = Inserter::new( + MemoryCatalogManager::new_with_table(DistTable::table(table.clone())), + partitions, + nodes, + mock_table_flownode_cache(1, vec![]).await, + true, + ) + .with_pending_rows_batcher(Some(batcher)); + let request = |value: i32| InsertRequest { + catalog_name: table.catalog_name.clone(), + schema_name: table.schema_name.clone(), + table_name: table.name.clone(), + columns_values: HashMap::from([ + ( + "a".to_string(), + Arc::new(Int32Vector::from_slice([value])) as VectorRef, + ), + ( + "ts".to_string(), + Arc::new(TimestampMillisecondVector::from_slice([ + 1000 + i64::from(value) + ])) as VectorRef, + ), + ]), + skip_wal, + }; + let (first, second) = tokio::join!( + inserter.handle_table_insert(request(1), influx_ctx), + inserter.handle_table_insert(request(2), opentsdb_ctx), + ); + ( + first.map(|output| match output.data { + OutputData::AffectedRows(rows) => rows, + _ => panic!("expected affected rows"), + }), + second.map(|output| match output.data { + OutputData::AffectedRows(rows) => rows, + _ => panic!("expected affected rows"), + }), + ) + } }) .await .expect("the row threshold did not dispatch the combined bulk insert"); @@ -474,7 +517,7 @@ mod tests { } #[tokio::test] - async fn test_submit_combines_then_routes_bulk_with_defaults() { + async fn test_inserter_counts_batched_rows_once() { for report_missing_row in [false, true] { assert_eq!(1, run_bulk_case(2, report_missing_row, false).await); } @@ -583,7 +626,7 @@ mod tests { FlowNotifier::new(cache.clone(), nodes.clone(), NonZeroUsize::new(16).unwrap()) .unwrap(); let inserter = Arc::new(Inserter::new( - catalog::memory::MemoryCatalogManager::new(), + MemoryCatalogManager::new(), partitions, nodes, cache, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index d6cb0cf13ec..f8a4de57654 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2306,6 +2306,15 @@ enable = true enable = true default_merge_mode = "last_non_null" +[pending_rows_batcher] +protocols = [] +pending_rows_flush_interval = "0s" +max_batch_rows = 100000 +max_concurrent_flushes = 256 +worker_channel_capacity = 65526 +max_inflight_requests = 3000 +flow_notification_queue_capacity = 1024 + [jaeger] enable = true