From 8d8ebd3cc5ae9c783c8b0425098bd89eb368f922 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 18 Sep 2026 06:09:20 +0000 Subject: [PATCH 01/63] perf: add flight coalesce regression case with high-cardinality aggregations (#9214) * perf: add flight coalesce regression case with high-cardinality aggregations Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf: add flight coalesce aggregations bench case Add a direct_readable_sst case exercising grouped aggregation over coalesced batches: 16 hosts x 4096 instances, 32 SSTs of 32768 rows, timestamp-major series layout, three SQL queries (aggregation, topk, count_by_host) each with a 10% max candidate latency regression threshold. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../flight_coalesce_aggregations/case.toml | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 tests/perf/query_cases/flight_coalesce_aggregations/case.toml diff --git a/tests/perf/query_cases/flight_coalesce_aggregations/case.toml b/tests/perf/query_cases/flight_coalesce_aggregations/case.toml new file mode 100644 index 00000000000..6cc20276a9f --- /dev/null +++ b/tests/perf/query_cases/flight_coalesce_aggregations/case.toml @@ -0,0 +1,101 @@ +# Qualification case for the Flight record-batch coalescing on the +# datanode-to-frontend path (#9167). mito2 emits small (2k-row) batches for +# timestamp-major layouts; the coalescer should merge them before they cross +# the internal Flight connection. High-cardinality tags make the aggregate and +# topk outputs wide, so the extra batches are visible end to end. + +[case] +name = "flight_coalesce_aggregations" +description = "Aggregate and topk queries over a high-cardinality direct-SST fixture exercising Flight batch coalescing" + +[scenario] +kind = "direct_readable_sst" +seed = 9167 + +[[scenario.tables]] +database = "public" +name = "flight_coalesce_metrics" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 4096, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(3)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +sst_count = 32 +rows_per_sst = 32768 +row_group_size = 8192 +series_count = 4096 +start_unix_nanos = 1_704_067_200_000_000_000 # 2024-01-01T00:00:00Z +step_nanos = 15_000_000_000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +[[scenario.queries]] +name = "aggregation" +kind = "sql" +query = ''' +SELECT host, instance, avg(value), max(value), count(*) +FROM flight_coalesce_metrics +GROUP BY host, instance +ORDER BY host, instance +''' +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "topk" +kind = "sql" +query = ''' +SELECT instance, max(value) +FROM flight_coalesce_metrics +GROUP BY instance +ORDER BY max(value) DESC +LIMIT 100 +''' +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "count_by_host" +kind = "sql" +query = ''' +SELECT host, count(*) +FROM flight_coalesce_metrics +GROUP BY host +''' +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 From 5ef8a46e3f7183ab8ca2795aaaddb3b7d14c024d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 18 Sep 2026 07:02:10 +0000 Subject: [PATCH 02/63] feat(function): add mergeable binary average states (#9062) * feat(function): add mergeable binary average states Signed-off-by: discord9 * feat(function): expose avg_calc as an OSS scalar Signed-off-by: discord9 * fix(function): borrow invalid AVG scalar argument type Signed-off-by: discord9 * test(function): verify public AVG state SQL finalization Signed-off-by: discord9 * fix(function): return canonical AVG1 state for empty window frames DataFusion's plain-aggregate window executor bypasses the accumulator and calls default_value() directly when a window frame contains no rows. create_udaf's SimpleAggregateUDF derives default_value from the return type, which yields SQL NULL for Binary output instead of the canonical AVG1 empty state, so empty frames and frames over only NULL inputs became observable differently (IS NULL, direct state saves, state comparisons). Register avg_state/avg_merge through a small AggregateUDFImpl (AvgUdaf) that keeps the existing accumulator and declares the canonical AVG1 empty state as default_value, restoring the documented contract. Add a unit test asserting default_value equals the empty accumulator's evaluate() and an sqlness case covering the empty-frame window. Signed-off-by: discord9 --------- Signed-off-by: discord9 Signed-off-by: discord9 Co-authored-by: discord9 --- src/common/function/src/aggrs/approximate.rs | 11 + .../function/src/aggrs/approximate/avg.rs | 594 ++++++++++++++++++ src/common/function/src/function_registry.rs | 2 + src/common/function/src/scalars.rs | 1 + src/common/function/src/scalars/avg_calc.rs | 266 ++++++++ .../standalone/common/aggregate/avg.result | 70 +++ .../cases/standalone/common/aggregate/avg.sql | 28 + .../cases/avg_state_binary/case.toml | 10 + .../cases/avg_state_binary/setup.sql | 14 + .../cases/avg_state_binary/verify.result | 36 ++ .../cases/avg_state_binary/verify.sql | 18 + 11 files changed, 1050 insertions(+) create mode 100644 src/common/function/src/aggrs/approximate/avg.rs create mode 100644 src/common/function/src/scalars/avg_calc.rs create mode 100644 tests/compatibility/cases/avg_state_binary/case.toml create mode 100644 tests/compatibility/cases/avg_state_binary/setup.sql create mode 100644 tests/compatibility/cases/avg_state_binary/verify.result create mode 100644 tests/compatibility/cases/avg_state_binary/verify.sql diff --git a/src/common/function/src/aggrs/approximate.rs b/src/common/function/src/aggrs/approximate.rs index bfab225755a..562723f6727 100644 --- a/src/common/function/src/aggrs/approximate.rs +++ b/src/common/function/src/aggrs/approximate.rs @@ -18,6 +18,7 @@ use datatypes::arrow::datatypes::DataType; use crate::aggrs::aggr_wrapper::DeltaMergeWrapper; use crate::function_registry::FunctionRegistry; +pub mod avg; pub mod hll; pub mod uddsketch; pub mod welford; @@ -26,6 +27,16 @@ pub(crate) struct ApproximateFunction; impl ApproximateFunction { pub fn register(registry: &FunctionRegistry) { + let avg_merge = avg::AvgAccumulator::merge_udf_impl(); + registry.register_aggr(avg::AvgAccumulator::state_udf_impl()); + registry.register_aggr(avg_merge.clone()); + registry.register_aggr(AggregateUDF::new_from_impl(DeltaMergeWrapper::new( + avg_merge.clone(), + avg::AVG_STATE_NAME, + vec![DataType::Binary], + DataType::Binary, + ))); + let uddsketch_state = uddsketch::UddSketchState::state_udf_impl(); let uddsketch_merge = uddsketch::UddSketchState::merge_udf_impl(); let uddsketch_delta = AggregateUDF::new_from_impl(DeltaMergeWrapper::new( diff --git a/src/common/function/src/aggrs/approximate/avg.rs b/src/common/function/src/aggrs/approximate/avg.rs new file mode 100644 index 00000000000..9ca952edb7c --- /dev/null +++ b/src/common/function/src/aggrs/approximate/avg.rs @@ -0,0 +1,594 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::arrow::array::{ArrayRef, Float64Array}; +use datafusion::arrow::compute::sum; +use datafusion::common::cast::{as_binary_array, as_primitive_array}; +use datafusion::common::not_impl_err; +use datafusion::error::{DataFusionError, Result as DfResult}; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::logical_expr::{ + Accumulator as DfAccumulator, AggregateUDF, AggregateUDFImpl, Signature, +}; +use datafusion_common::ScalarValue; +use datatypes::arrow::datatypes::{DataType, Float64Type}; + +pub const AVG_STATE_NAME: &str = "avg_state"; +pub const AVG_MERGE_NAME: &str = "avg_merge"; + +const ENCODED_LEN: usize = 20; +const MAGIC: &[u8; 4] = b"AVG1"; + +/// The portable state used by the Float64 average aggregate functions. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct AvgState { + count: u64, + sum: f64, +} + +impl Default for AvgState { + fn default() -> Self { + Self { count: 0, sum: 0.0 } + } +} + +impl AvgState { + /// Returns the exact AVG1 representation of this state. + pub(crate) fn encode(&self) -> [u8; ENCODED_LEN] { + let mut encoded = [0; ENCODED_LEN]; + encoded[..4].copy_from_slice(MAGIC); + encoded[4..12].copy_from_slice(&self.count.to_le_bytes()); + encoded[12..20].copy_from_slice(&self.sum.to_bits().to_le_bytes()); + encoded + } + + /// Decodes and validates an AVG1 state. + pub fn decode(encoded: &[u8]) -> DfResult { + if encoded.len() != ENCODED_LEN || &encoded[..4] != MAGIC { + return Err(invalid_state()); + } + let count = decode_u64(encoded, 4); + let sum = f64::from_bits(decode_u64(encoded, 12)); + if count == 0 && sum.to_bits() != 0 { + return Err(invalid_state()); + } + Ok(Self { count, sum }) + } + + /// Returns the number of non-null input values in this state. + pub(crate) fn count(&self) -> u64 { + self.count + } + + /// Returns the average, or `None` for the canonical empty state. + pub fn average(&self) -> Option { + (self.count() != 0).then(|| self.sum / self.count() as f64) + } +} + +fn decode_u64(encoded: &[u8], offset: usize) -> u64 { + let mut bytes = [0; 8]; + bytes.copy_from_slice(&encoded[offset..offset + 8]); + u64::from_le_bytes(bytes) +} + +fn invalid_state() -> DataFusionError { + DataFusionError::Execution("Invalid AVG1 state".to_string()) +} + +fn count_overflow() -> DataFusionError { + DataFusionError::Execution("AVG count overflow".to_string()) +} + +/// The `avg_state` / `avg_merge` aggregate UDF. +/// +/// Declares a canonical AVG1 empty state as `default_value` so window frames +/// without rows observe the same contract as the accumulator's `evaluate`. +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +struct AvgUdaf { + name: &'static str, + signature: Signature, + input: InputKind, +} + +impl AggregateUDFImpl for AvgUdaf { + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Binary) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> DfResult> { + if acc_args.is_distinct { + return not_impl_err!("AVG DISTINCT aggregations are not available"); + } + let input = match acc_args.exprs[0].data_type(acc_args.schema)? { + DataType::Float64 => InputKind::Float64, + DataType::Binary => InputKind::Binary, + data_type => return not_impl_err!("AVG functions do not support {data_type:?}"), + }; + Ok(Box::new(AvgAccumulator { + state: AvgState::default(), + input, + })) + } + + fn default_value(&self, _data_type: &DataType) -> DfResult { + Ok(ScalarValue::Binary(Some( + AvgState::default().encode().to_vec(), + ))) + } +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)] +enum InputKind { + Float64, + Binary, +} + +/// Accumulates and merges AVG1 states. +#[derive(Debug)] +pub(crate) struct AvgAccumulator { + state: AvgState, + input: InputKind, +} + +impl Default for AvgAccumulator { + fn default() -> Self { + Self { + state: AvgState::default(), + input: InputKind::Float64, + } + } +} + +impl AvgAccumulator { + pub fn state_udf_impl() -> AggregateUDF { + AggregateUDF::new_from_impl(AvgUdaf { + name: AVG_STATE_NAME, + signature: Signature::exact( + vec![DataType::Float64], + datafusion::logical_expr::Volatility::Immutable, + ), + input: InputKind::Float64, + }) + } + + pub fn merge_udf_impl() -> AggregateUDF { + AggregateUDF::new_from_impl(AvgUdaf { + name: AVG_MERGE_NAME, + signature: Signature::exact( + vec![DataType::Binary], + datafusion::logical_expr::Volatility::Immutable, + ), + input: InputKind::Binary, + }) + } + + fn update_float64(&mut self, array: &ArrayRef) -> DfResult<()> { + let array = as_primitive_array::(array)?; + let mut count = self.state.count; + for _ in array.iter().flatten() { + count = count.checked_add(1).ok_or_else(count_overflow)?; + } + let sum = sum(array) + .map(|batch_sum| self.state.sum + batch_sum) + .unwrap_or(self.state.sum); + self.state = AvgState { count, sum }; + Ok(()) + } + + fn merge_states(&mut self, array: &ArrayRef) -> DfResult<()> { + let array = as_binary_array(array)?; + let states = array + .iter() + .flatten() + .map(AvgState::decode) + .collect::>>()?; + let count = states.iter().try_fold(self.state.count, |count, state| { + count.checked_add(state.count).ok_or_else(count_overflow) + })?; + let sums = states + .iter() + .filter(|state| state.count != 0) + .map(|state| Some(state.sum)) + .collect::>(); + let sum = sum(&Float64Array::from(sums)) + .map(|batch_sum| self.state.sum + batch_sum) + .unwrap_or(self.state.sum); + self.state = AvgState { count, sum }; + Ok(()) + } +} + +impl DfAccumulator for AvgAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> { + let array = &values[0]; + match (self.input, array.data_type()) { + (InputKind::Float64, DataType::Float64) => self.update_float64(array), + (InputKind::Binary, DataType::Binary) => self.merge_states(array), + (_, data_type) => not_impl_err!("AVG input type does not match: {data_type:?}"), + } + } + + fn evaluate(&mut self) -> DfResult { + Ok(ScalarValue::Binary(Some(self.state.encode().to_vec()))) + } + + fn size(&self) -> usize { + std::mem::size_of::() + } + + fn state(&mut self) -> DfResult> { + Ok(vec![ScalarValue::Binary(Some( + self.state.encode().to_vec(), + ))]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> { + self.merge_states(&states[0]) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{BinaryArray, Float64Array}; + use datafusion_common::ScalarValue; + use datafusion_common::arrow::datatypes::DataType; + use datafusion_expr::TypeSignature; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::{Column, lit as physical_lit}; + + use super::*; + use crate::aggrs::aggr_wrapper::{aggr_delta_merge_func_name, aggr_state_func_name}; + use crate::function_registry::FUNCTION_REGISTRY; + + fn state(count: u64, sum: f64) -> Vec { + AvgState { count, sum }.encode().to_vec() + } + + #[test] + fn codec_golden_and_roundtrip() { + let empty = AvgState::default().encode(); + assert_eq!(empty.len(), ENCODED_LEN); + assert_eq!(empty.as_slice(), b"AVG1\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"); + let mut accumulator = AvgAccumulator::default(); + accumulator + .update_batch(&[Arc::new(Float64Array::from(vec![Some(1.5)]))]) + .unwrap(); + let one = accumulator.state.encode(); + assert_eq!(one.len(), ENCODED_LEN); + assert_eq!( + one.as_slice(), + &[ + b'A', b'V', b'G', b'1', 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xf8, 0x3f, + ] + ); + assert_eq!(&one[12..20], &1.5f64.to_bits().to_le_bytes()); + assert_eq!(AvgState::decode(&empty).unwrap().encode(), empty); + assert_eq!(AvgState::decode(&one).unwrap().encode(), one); + } + + #[test] + fn codec_rejects_malformed_states() { + assert!(AvgState::decode(b"").is_err()); + assert!(AvgState::decode(&[0; 19]).is_err()); + assert!(AvgState::decode(&[0; 21]).is_err()); + let mut avg2 = AvgState::default().encode(); + avg2[..4].copy_from_slice(b"AVG2"); + assert!(AvgState::decode(&avg2).is_err()); + let mut wrong_magic = AvgState::default().encode(); + wrong_magic[0] = b'X'; + assert!(AvgState::decode(&wrong_magic).is_err()); + for sum in [1.0, -0.0] { + assert!(AvgState::decode(&state(0, sum)).is_err()); + } + let mut count = AvgState { + count: 0x0102_0304_0506_0708, + sum: 0.0, + } + .encode(); + assert_eq!(&count[4..12], &0x0102_0304_0506_0708u64.to_le_bytes()); + count[4..12].reverse(); + assert_ne!( + AvgState::decode(&count).unwrap().count(), + 0x0102_0304_0506_0708 + ); + let mut sum = AvgState { count: 1, sum: 1.5 }.encode(); + assert_eq!(&sum[12..20], &1.5f64.to_bits().to_le_bytes()); + sum[12..20].reverse(); + assert_ne!(AvgState::decode(&sum).unwrap().average(), Some(1.5)); + } + + #[test] + fn codec_preserves_populated_float_bits() { + for bits in [ + 0.0f64.to_bits(), + (-0.0f64).to_bits(), + f64::INFINITY.to_bits(), + f64::NEG_INFINITY.to_bits(), + 0x7ff8_0000_0000_0001, + 0x7ff0_0000_0000_0001, + ] { + let encoded = state(1, f64::from_bits(bits)); + assert_eq!( + AvgState::decode(&encoded).unwrap().encode().as_slice(), + encoded + ); + } + } + + #[test] + fn distinct_is_rejected() { + let udf = AvgAccumulator::state_udf_impl(); + let schema = arrow_schema::Schema::empty(); + let expr = physical_lit(1.0f64); + let field = Arc::new(arrow_schema::Field::new("in", DataType::Float64, true)); + let args = AccumulatorArgs { + return_field: Arc::new(arrow_schema::Field::new("out", DataType::Binary, true)), + schema: &schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: AVG_STATE_NAME, + is_distinct: true, + exprs: std::slice::from_ref(&expr), + expr_fields: std::slice::from_ref(&field), + }; + assert!(udf.accumulator(args).is_err()); + } + + #[test] + fn state_counts_nulls_and_empty_is_canonical() { + let mut accumulator = AvgAccumulator::default(); + accumulator + .update_batch(&[Arc::new(Float64Array::from(vec![None, None]))]) + .unwrap(); + assert_eq!(accumulator.state.encode(), AvgState::default().encode()); + accumulator + .update_batch(&[Arc::new(Float64Array::from(vec![ + Some(1.0), + None, + Some(3.0), + Some(8.0), + ]))]) + .unwrap(); + assert_eq!(accumulator.state.count(), 3); + assert_eq!(accumulator.state.average(), Some(4.0)); + } + + #[test] + fn default_value_matches_empty_accumulator_evaluate() { + for udf in [ + AvgAccumulator::state_udf_impl(), + AvgAccumulator::merge_udf_impl(), + ] { + let default = udf.default_value(&DataType::Binary).unwrap(); + let expected = ScalarValue::Binary(Some(AvgState::default().encode().to_vec())); + assert_eq!(default, expected); + let mut empty = AvgAccumulator::default(); + assert_eq!( + default, + empty.evaluate().unwrap(), + "{} default_value must equal the empty accumulator evaluate", + udf.name() + ); + } + } + + #[test] + fn merge_preserves_populated_negative_zero_for_empty_input() { + let mut accumulator = AvgAccumulator { + state: AvgState { + count: 1, + sum: -0.0, + }, + input: InputKind::Binary, + }; + let expected = accumulator.state.encode(); + accumulator + .update_batch(&[Arc::new(BinaryArray::from(vec![ + None, + Some(AvgState::default().encode().as_slice()), + ]))]) + .unwrap(); + assert_eq!(accumulator.state.encode(), expected); + } + + #[test] + fn merge_ignores_nulls_and_merges_weighted_states() { + let mut accumulator = AvgAccumulator { + state: AvgState::default(), + input: InputKind::Binary, + }; + accumulator + .update_batch(&[Arc::new(BinaryArray::from(vec![ + Some(state(2, 4.0).as_slice()), + None, + Some(state(3, 15.0).as_slice()), + ]))]) + .unwrap(); + assert_eq!(accumulator.state.count(), 5); + assert_eq!(accumulator.state.average(), Some(19.0 / 5.0)); + let before = accumulator.state; + assert!( + accumulator + .update_batch(&[Arc::new(BinaryArray::from(vec![Some(&[][..])]))]) + .is_err() + ); + assert_eq!(accumulator.state, before); + } + + #[test] + fn overflow_does_not_mutate_update_or_merge() { + let mut update = AvgAccumulator { + state: AvgState { + count: u64::MAX, + sum: 1.0, + }, + input: InputKind::Float64, + }; + let before = update.state; + assert!( + update + .update_batch(&[Arc::new(Float64Array::from(vec![Some(2.0)]))]) + .is_err() + ); + assert_eq!(update.state, before); + + let mut merge = AvgAccumulator { + state: AvgState { + count: u64::MAX, + sum: 1.0, + }, + input: InputKind::Binary, + }; + let before = merge.state; + assert!( + merge + .update_batch(&[Arc::new(BinaryArray::from(vec![Some( + state(1, 2.0).as_slice() + )]))]) + .is_err() + ); + assert_eq!(merge.state, before); + } + + #[test] + fn registered_delta_merge_has_four_way_and_malformed_behavior() { + let udf = FUNCTION_REGISTRY + .get_aggr_func(&aggr_delta_merge_func_name(AVG_STATE_NAME)) + .unwrap(); + assert_eq!(udf.name(), "__avg_state_delta_merge"); + assert_eq!( + udf.signature().type_signature, + TypeSignature::Exact(vec![DataType::Binary, DataType::Binary]) + ); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("delta", DataType::Binary, true), + arrow_schema::Field::new("persisted", DataType::Binary, true), + ])); + let expr = AggregateExprBuilder::new( + Arc::new(udf), + vec![ + Arc::new(Column::new("delta", 0)), + Arc::new(Column::new("persisted", 1)), + ], + ) + .schema(schema) + .alias("avg_delta_merge") + .build() + .unwrap(); + let delta = state(2, 3.0); + let persisted = state(2, 7.0); + for (left, right, expected) in [ + ( + Some(delta.as_slice()), + None, + AvgState { count: 2, sum: 3.0 }.encode(), + ), + ( + None, + Some(persisted.as_slice()), + AvgState { count: 2, sum: 7.0 }.encode(), + ), + (None, None, AvgState::default().encode()), + ( + Some(delta.as_slice()), + Some(persisted.as_slice()), + AvgState { + count: 4, + sum: 10.0, + } + .encode(), + ), + ] { + let mut accumulator = expr.create_accumulator().unwrap(); + accumulator + .update_batch(&[ + Arc::new(BinaryArray::from(vec![left])), + Arc::new(BinaryArray::from(vec![right])), + ]) + .unwrap(); + let ScalarValue::Binary(Some(actual)) = accumulator.evaluate().unwrap() else { + panic!("AVG delta merge state must be binary"); + }; + assert_eq!(actual.as_slice(), expected.as_slice()); + } + let mut accumulator = expr.create_accumulator().unwrap(); + assert!( + accumulator + .update_batch(&[ + Arc::new(BinaryArray::from(vec![Some(&[][..])])), + Arc::new(BinaryArray::from(vec![None])), + ]) + .is_err() + ); + let mut accumulator = expr.create_accumulator().unwrap(); + assert!( + accumulator + .update_batch(&[ + Arc::new(BinaryArray::from(vec![None])), + Arc::new(BinaryArray::from(vec![Some(&[][..])])), + ]) + .is_err() + ); + let mut accumulator = expr.create_accumulator().unwrap(); + accumulator + .update_batch(&[ + Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])), + Arc::new(BinaryArray::from(vec![None])), + ]) + .unwrap(); + assert!( + accumulator + .update_batch(&[ + Arc::new(BinaryArray::from(vec![Some(delta.as_slice())])), + Arc::new(BinaryArray::from(vec![Some( + AvgState { + count: u64::MAX, + sum: 1.0 + } + .encode() + .as_slice() + )])), + ]) + .is_err() + ); + } + + #[test] + fn avg_registry_does_not_replace_native_state_registry() { + let avg = FUNCTION_REGISTRY.get_aggr_func(AVG_STATE_NAME).unwrap(); + let native = FUNCTION_REGISTRY + .get_aggr_func(&aggr_state_func_name("avg")) + .unwrap(); + assert_eq!( + avg.return_type(&[DataType::Float64]).unwrap(), + DataType::Binary + ); + assert!(matches!( + native.return_type(&[DataType::Float64]).unwrap(), + DataType::Struct(_) + )); + } +} diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 80c34eeb91f..c8e9874abc9 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -29,6 +29,7 @@ use crate::aggrs::vector::VectorFunction as VectorAggrFunction; use crate::function::{Function, FunctionRef}; use crate::function_factory::ScalarFunctionFactory; use crate::scalars::anomaly::AnomalyFunction; +use crate::scalars::avg_calc::AvgCalcFunction; use crate::scalars::date::DateFunction; use crate::scalars::expression::ExpressionFunction; use crate::scalars::hll_count::HllCalcFunction; @@ -215,6 +216,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| TimestampFunction::register(&function_registry); DateFunction::register(&function_registry); ExpressionFunction::register(&function_registry); + AvgCalcFunction::register(&function_registry); UddSketchCalcFunction::register(&function_registry); UddSketchRankFunction::register(&function_registry); HllCalcFunction::register(&function_registry); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 1a8c22e2d6f..306d1cd17e9 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -25,6 +25,7 @@ pub mod primary_key; pub(crate) mod string; pub mod vector; +pub(crate) mod avg_calc; pub(crate) mod hll_count; pub mod ip; #[cfg(test)] diff --git a/src/common/function/src/scalars/avg_calc.rs b/src/common/function/src/scalars/avg_calc.rs new file mode 100644 index 00000000000..cae76cb2b1f --- /dev/null +++ b/src/common/function/src/scalars/avg_calc.rs @@ -0,0 +1,266 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Implementation of the scalar function `avg_calc`. + +use std::fmt; +use std::fmt::Display; +use std::sync::Arc; + +use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; +use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; +use datatypes::arrow::datatypes::DataType; + +use crate::aggrs::approximate::avg::AvgState; +use crate::function::Function; +use crate::function_registry::FunctionRegistry; + +const NAME: &str = "avg_calc"; + +/// Calculates an average from a serialized AVG1 state. +#[derive(Debug)] +pub(crate) struct AvgCalcFunction { + signature: Signature, +} + +impl AvgCalcFunction { + pub fn register(registry: &FunctionRegistry) { + registry.register_scalar(Self::default()); + } +} + +impl Default for AvgCalcFunction { + fn default() -> Self { + Self { + signature: Signature::exact(vec![DataType::Binary], Volatility::Immutable), + } + } +} + +impl Display for AvgCalcFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", NAME.to_ascii_uppercase()) + } +} + +impl Function for AvgCalcFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Float64) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [arg] = datafusion_common::utils::take_function_args(self.name(), &args.args)?; + match arg { + ColumnarValue::Scalar(ScalarValue::Binary(state)) => { + Ok(ColumnarValue::Scalar(ScalarValue::Float64( + state + .as_deref() + .map(AvgState::decode) + .transpose()? + .and_then(|state| state.average()), + ))) + } + ColumnarValue::Scalar(ScalarValue::Null) => { + Ok(ColumnarValue::Scalar(ScalarValue::Float64(None))) + } + ColumnarValue::Array(states) => { + let Some(states) = states.as_binary_opt::() else { + return Err(invalid_type(self.name(), states.data_type())); + }; + let mut builder = Float64Builder::with_capacity(states.len()); + for state in states.iter() { + builder.append_option(match state { + Some(state) => AvgState::decode(state)?.average(), + None => None, + }); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } + _ => Err(invalid_type(self.name(), &arg.data_type())), + } + } +} + +fn invalid_type(name: &str, data_type: &DataType) -> DataFusionError { + DataFusionError::Execution(format!( + "'{name}' expects argument to be Binary datatype, got {data_type}" + )) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion::arrow::array::{Array, AsArray, BinaryArray, Float64Array}; + use datafusion::logical_expr::Accumulator; + use datafusion::prelude::SessionContext; + use datafusion_common::arrow::datatypes::Float64Type; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; + + use super::*; + use crate::aggrs::approximate::avg::AvgAccumulator; + use crate::function::{Function, FunctionContext}; + use crate::function_registry::FUNCTION_REGISTRY; + + fn produce_state(values: Vec>) -> Vec { + let mut accumulator = AvgAccumulator::default(); + accumulator + .update_batch(&[Arc::new(Float64Array::from(values))]) + .unwrap(); + let ScalarValue::Binary(Some(state)) = accumulator.evaluate().unwrap() else { + panic!("AVG state must be binary"); + }; + state + } + + fn invoke(arg: ColumnarValue, number_rows: usize) -> datafusion_common::Result { + AvgCalcFunction::default().invoke_with_args(ScalarFunctionArgs { + args: vec![arg], + arg_fields: vec![], + number_rows, + return_field: Arc::new(Field::new("x", DataType::Float64, true)), + config_options: Arc::new(Default::default()), + }) + } + + #[test] + fn scalar_and_array_states_decode_to_averages() { + let state = produce_state(vec![Some(1.0), Some(2.0), Some(6.0)]); + let ColumnarValue::Scalar(ScalarValue::Float64(Some(value))) = + invoke(ColumnarValue::Scalar(ScalarValue::Binary(Some(state))), 1).unwrap() + else { + panic!("Expected Float64 scalar"); + }; + assert_eq!(value, 3.0); + + let ColumnarValue::Scalar(ScalarValue::Float64(None)) = + invoke(ColumnarValue::Scalar(ScalarValue::Binary(None)), 1).unwrap() + else { + panic!("Expected NULL Float64 scalar"); + }; + let empty = produce_state(vec![None]); + let ColumnarValue::Scalar(ScalarValue::Float64(None)) = invoke( + ColumnarValue::Scalar(ScalarValue::Binary(Some(empty.clone()))), + 1, + ) + .unwrap() else { + panic!("Expected NULL Float64 scalar"); + }; + + let infinity = produce_state(vec![Some(f64::INFINITY)]); + let nan = produce_state(vec![Some(f64::NAN)]); + let ColumnarValue::Array(result) = invoke( + ColumnarValue::Array(Arc::new(BinaryArray::from(vec![ + Some(empty.as_slice()), + None, + Some(infinity.as_slice()), + Some(nan.as_slice()), + ]))), + 4, + ) + .unwrap() else { + panic!("Expected Float64 array"); + }; + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert!(result.is_null(1)); + assert_eq!(result.value(2), f64::INFINITY); + assert!(result.value(3).is_nan()); + } + + #[test] + fn malformed_and_unknown_version_states_fail_the_whole_batch() { + let valid = produce_state(vec![Some(3.0)]); + let mut unknown_version = valid.clone(); + unknown_version[..4].copy_from_slice(b"AVG2"); + + assert!( + invoke( + ColumnarValue::Scalar(ScalarValue::Binary(Some(unknown_version.clone()))), + 1, + ) + .is_err() + ); + assert!( + invoke( + ColumnarValue::Array(Arc::new(BinaryArray::from(vec![ + Some(valid.as_slice()), + Some(b"malformed".as_slice()), + Some(unknown_version.as_slice()), + ]))), + 3, + ) + .is_err() + ); + } + + #[tokio::test] + async fn registry_query_decodes_avg_state_and_weighted_avg_merge() { + let ctx = SessionContext::new(); + let avg_calc = FUNCTION_REGISTRY + .get_function(NAME) + .expect("avg_calc must be registered") + .provide(FunctionContext::default()); + ctx.register_udf(avg_calc); + for name in ["avg_state", "avg_merge"] { + ctx.register_udaf( + FUNCTION_REGISTRY + .get_aggr_func(name) + .expect("AVG aggregate must be registered"), + ); + } + + let batches = ctx + .sql( + "SELECT avg_calc(avg_state(CAST(value AS DOUBLE))) FROM \ + (VALUES (1.0), (2.0), (6.0)) AS values_table(value)", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + let result = batches[0].column(0).as_primitive::(); + assert_eq!(result.value(0), 3.0); + + let batches = ctx + .sql( + "WITH states AS (\ + SELECT avg_state(CAST(value AS DOUBLE)) AS state FROM (VALUES (1.0), (3.0)) AS left_values(value) \ + UNION ALL \ + SELECT avg_state(CAST(value AS DOUBLE)) AS state FROM (VALUES (6.0)) AS right_values(value)\ + ) SELECT avg_calc(avg_merge(state)) FROM states", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + let result = batches[0].column(0).as_primitive::(); + assert_eq!(result.value(0), 10.0 / 3.0); + } +} diff --git a/tests/cases/standalone/common/aggregate/avg.result b/tests/cases/standalone/common/aggregate/avg.result index 99fd8a09e4c..686c7df9e5b 100644 --- a/tests/cases/standalone/common/aggregate/avg.result +++ b/tests/cases/standalone/common/aggregate/avg.result @@ -89,6 +89,76 @@ SELECT AVG(i), AVG(j) FROM vals; -- FIXME(dennis): AVG(DISTINCT) not supported -- https://github.com/apache/datafusion/issues/2408 -- SELECT AVG(DISTINCT i), AVG(DISTINCT j) FROM vals; +-- AVG1 state scalar calculation +SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers; + ++---------------------------------+ +| avg_calc(avg_state(integers.i)) | ++---------------------------------+ +| 2.0 | ++---------------------------------+ + +-- Merged states retain their unequal group weights. +WITH states AS ( + SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i < 3 + UNION ALL + SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i >= 3 +) +SELECT avg_calc(avg_merge(state)) FROM states; + ++-----------------------------------+ +| avg_calc(avg_merge(states.state)) | ++-----------------------------------+ +| 2.0 | ++-----------------------------------+ + +-- Empty, null-only, and null binary states calculate to NULL. +SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers WHERE i > 100; + ++---------------------------------+ +| avg_calc(avg_state(integers.i)) | ++---------------------------------+ +| | ++---------------------------------+ + +SELECT avg_calc(avg_state(NULL::DOUBLE)); + ++---------------------------+ +| avg_calc(avg_state(NULL)) | ++---------------------------+ +| | ++---------------------------+ + +SELECT avg_calc(NULL::BYTEA); + ++----------------+ +| avg_calc(NULL) | ++----------------+ +| | ++----------------+ + +-- Empty window frames return the canonical AVG1 state, not SQL NULL. +SELECT + i, + avg_state(CAST(i AS DOUBLE)) OVER ( + ORDER BY i + ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING + ) AS state +FROM integers; + ++---+------------------------------------------+ +| i | state | ++---+------------------------------------------+ +| 1 | 4156473100000000000000000000000000000000 | +| 2 | 415647310100000000000000000000000000f03f | +| 3 | 4156473102000000000000000000000000000840 | ++---+------------------------------------------+ + +-- Invalid AVG1 state propagates an error. +SELECT avg_calc(X'00'); + +Error: 3001(EngineExecuteQuery), Execution error: Invalid AVG1 state + -- cleanup DROP TABLE integers; diff --git a/tests/cases/standalone/common/aggregate/avg.sql b/tests/cases/standalone/common/aggregate/avg.sql index cbb12edcbee..0667bbcfde8 100644 --- a/tests/cases/standalone/common/aggregate/avg.sql +++ b/tests/cases/standalone/common/aggregate/avg.sql @@ -40,6 +40,34 @@ SELECT AVG(i), AVG(j) FROM vals; -- https://github.com/apache/datafusion/issues/2408 -- SELECT AVG(DISTINCT i), AVG(DISTINCT j) FROM vals; +-- AVG1 state scalar calculation +SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers; + +-- Merged states retain their unequal group weights. +WITH states AS ( + SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i < 3 + UNION ALL + SELECT avg_state(CAST(i AS DOUBLE)) AS state FROM integers WHERE i >= 3 +) +SELECT avg_calc(avg_merge(state)) FROM states; + +-- Empty, null-only, and null binary states calculate to NULL. +SELECT avg_calc(avg_state(CAST(i AS DOUBLE))) FROM integers WHERE i > 100; +SELECT avg_calc(avg_state(NULL::DOUBLE)); +SELECT avg_calc(NULL::BYTEA); + +-- Empty window frames return the canonical AVG1 state, not SQL NULL. +SELECT + i, + avg_state(CAST(i AS DOUBLE)) OVER ( + ORDER BY i + ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING + ) AS state +FROM integers; + +-- Invalid AVG1 state propagates an error. +SELECT avg_calc(X'00'); + -- cleanup DROP TABLE integers; diff --git a/tests/compatibility/cases/avg_state_binary/case.toml b/tests/compatibility/cases/avg_state_binary/case.toml new file mode 100644 index 00000000000..0049fca8c39 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/case.toml @@ -0,0 +1,10 @@ +name = "avg_state_binary" +reason = "Verify persisted AVG1 binary states are decoded, merged, and reproduced exactly after upgrade." +introduced_by = "PR #9035" +topologies = ["distributed", "standalone"] +from_range = ["*"] +# The runner compares core versions only, so this also admits other 1.3.0 prereleases. +to_range = [">=v1.3.0-alpha.1"] +features = ["table", "query", "aggregate"] +owner = "query" +namespace = "avg_state_binary" diff --git a/tests/compatibility/cases/avg_state_binary/setup.sql b/tests/compatibility/cases/avg_state_binary/setup.sql new file mode 100644 index 00000000000..2bb2db577ee --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/setup.sql @@ -0,0 +1,14 @@ +CREATE TABLE avg1_states ( + seq_id INT PRIMARY KEY, + state BINARY, + ts TIMESTAMP TIME INDEX +); + +-- AVG1: magic (4 bytes), little-endian u64 count, little-endian f64 sum. +INSERT INTO avg1_states (seq_id, state, ts) VALUES + (1, X'4156473102000000000000000000000000000840', '2026-01-01 00:00:00'), + (2, X'4156473101000000000000000000000000001840', '2026-01-01 00:00:01'), + (3, NULL, '2026-01-01 00:00:02'), + (4, X'4156473100000000000000000000000000000000', '2026-01-01 00:00:03'); + +ADMIN FLUSH_TABLE('avg1_states'); diff --git a/tests/compatibility/cases/avg_state_binary/verify.result b/tests/compatibility/cases/avg_state_binary/verify.result new file mode 100644 index 00000000000..c3eba781a49 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/verify.result @@ -0,0 +1,36 @@ +-- The persisted states have counts 2 and 1 and sums 3.0 and 6.0. +-- Null and canonical empty states do not change the merged AVG1 state. +SELECT avg_merge(state) = X'4156473103000000000000000000000000002240' AS merged_state_matches +FROM avg1_states; + ++----------------------+ +| merged_state_matches | ++----------------------+ +| true | ++----------------------+ + +-- A new state over the equivalent Float64 values has the exact same AVG1 bytes. +WITH generated_values AS ( + SELECT CAST(1.0 AS DOUBLE) AS value + UNION ALL SELECT CAST(2.0 AS DOUBLE) + UNION ALL SELECT CAST(6.0 AS DOUBLE) +) +SELECT avg_state(value) = X'4156473103000000000000000000000000002240' AS generated_state_matches +FROM generated_values; + ++-------------------------+ +| generated_state_matches | ++-------------------------+ +| true | ++-------------------------+ + +-- A null-only merge is the canonical empty AVG1 state. +SELECT avg_merge(state) = X'4156473100000000000000000000000000000000' AS null_state_is_empty +FROM avg1_states +WHERE state IS NULL; + ++---------------------+ +| null_state_is_empty | ++---------------------+ +| true | ++---------------------+ diff --git a/tests/compatibility/cases/avg_state_binary/verify.sql b/tests/compatibility/cases/avg_state_binary/verify.sql new file mode 100644 index 00000000000..30cade275b6 --- /dev/null +++ b/tests/compatibility/cases/avg_state_binary/verify.sql @@ -0,0 +1,18 @@ +-- The persisted states have counts 2 and 1 and sums 3.0 and 6.0. +-- Null and canonical empty states do not change the merged AVG1 state. +SELECT avg_merge(state) = X'4156473103000000000000000000000000002240' AS merged_state_matches +FROM avg1_states; + +-- A new state over the equivalent Float64 values has the exact same AVG1 bytes. +WITH generated_values AS ( + SELECT CAST(1.0 AS DOUBLE) AS value + UNION ALL SELECT CAST(2.0 AS DOUBLE) + UNION ALL SELECT CAST(6.0 AS DOUBLE) +) +SELECT avg_state(value) = X'4156473103000000000000000000000000002240' AS generated_state_matches +FROM generated_values; + +-- A null-only merge is the canonical empty AVG1 state. +SELECT avg_merge(state) = X'4156473100000000000000000000000000000000' AS null_state_is_empty +FROM avg1_states +WHERE state IS NULL; From 5b631fd8eb8fc99e1b4dc1321c0e1850b6924fe3 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:08:35 +0000 Subject: [PATCH 03/63] feat(json2)!: remove nullable and default type hint options (#9213) feat(json2): remove nullable and default type hint options --- src/datatypes/src/json.rs | 107 ++---------- src/datatypes/src/vectors/json/array.rs | 2 - src/datatypes/src/vectors/json/builder.rs | 10 -- src/mito2/src/compaction/json2.rs | 18 -- src/mito2/src/memtable/bulk.rs | 2 - src/pipeline/src/etl/transform.rs | 73 +------- .../src/etl/transform/transformer/greptime.rs | 4 - .../transform/transformer/greptime/coerce.rs | 2 - src/pipeline/tests/json_parse.rs | 1 - src/sql/src/parsers/alter_parser.rs | 6 +- src/sql/src/parsers/create_parser/json.rs | 165 ++++-------------- src/sql/src/statements.rs | 2 - src/sql/src/statements/create.rs | 162 +++-------------- .../standalone/common/types/json/json2.result | 2 +- .../common/types/json/json2_alter.result | 8 +- .../common/types/json/json2_alter.sql | 8 +- .../common/types/json/json2_type_hints.result | 96 +++------- .../common/types/json/json2_type_hints.sql | 40 +---- 18 files changed, 112 insertions(+), 596 deletions(-) diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index e156e387330..a5b4940e048 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -31,7 +31,6 @@ use snafu::ResultExt; use crate::data_type::ConcreteDataType; use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; -use crate::schema::ColumnDefaultConstraint; use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::value::{ListValue, StructValue, Value}; @@ -91,9 +90,6 @@ pub struct JsonTypeHint { pub path: Vec, #[serde(rename = "type")] pub data_type: ConcreteDataType, - pub nullable: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub default_constraint: Option, pub inverted_index: bool, } @@ -217,20 +213,6 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { .fail(); } }; - let non_finite_default = match &hint.default_constraint { - Some(ColumnDefaultConstraint::Value(Value::Float32(value))) => !value.0.is_finite(), - Some(ColumnDefaultConstraint::Value(Value::Float64(value))) => !value.0.is_finite(), - _ => false, - }; - if non_finite_default { - return InvalidJson2SettingsSnafu { - reason: format!( - "JSON2 type hint default for '{}' must be finite", - hint.path.join(".") - ), - } - .fail(); - } validate_type_hint(&mut object, &hint.path, data_type)?; } Ok(()) @@ -319,12 +301,12 @@ fn encode_json_object_with_context<'a>( object.insert(key, value.into_variant()); } - apply_missing_type_hints(&mut object, context)?; + fill_missing_type_hints(&mut object, context)?; Ok(JsonValue::new(JsonVariant::Object(object))) } -fn apply_missing_type_hints( +fn fill_missing_type_hints( object: &mut BTreeMap, context: &mut JsonContext, ) -> Result<()> { @@ -351,8 +333,7 @@ fn insert_missing_type_hint( if is_leaf { if !object.contains_key(key) { - let value = encode_missing_type_hint_value(hint, field_context)?; - object.insert(key.clone(), value.into_variant()); + object.insert(key.clone(), JsonValue::null().into_variant()); } return Ok(()); } @@ -380,46 +361,13 @@ fn insert_missing_type_hint( } } -fn encode_missing_type_hint_value( - hint: &JsonTypeHint, - context: &mut JsonContext, -) -> Result { - if let Some(default_constraint) = &hint.default_constraint { - let value = default_constraint.create_default(&hint.data_type, hint.nullable)?; - let json = decode_primitive_value(value)?; - return encode_json_value_with_hint(json, hint, context); - } - - if hint.nullable { - Ok(JsonValue::null()) - } else { - error::InvalidJsonSnafu { - value: format!( - "missing non-null JSON2 type hint path {}", - hint.path.join(".") - ), - } - .fail() - } -} - fn encode_json_value_with_hint( json: Json, hint: &JsonTypeHint, context: &mut JsonContext, ) -> Result { if json.is_null() { - return if hint.nullable { - Ok(JsonValue::null()) - } else { - error::InvalidJsonSnafu { - value: format!( - "JSON2 type hint path {} is not nullable", - context.path.join(".") - ), - } - .fail() - }; + return Ok(JsonValue::null()); } let invalid_type = || { @@ -650,7 +598,7 @@ mod tests { } #[test] - fn test_json_settings_forward_compatibility() + fn test_json_settings_deserializes_legacy_type_hint_constraints() -> std::result::Result<(), Box> { let json_str = r#"{ "type_hints": [ @@ -687,15 +635,11 @@ mod tests { JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(18))), inverted_index: true, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ], @@ -711,15 +655,11 @@ mod tests { JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(18))), inverted_index: true, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ]; @@ -758,8 +698,6 @@ mod tests { let hint = |path, data_type| JsonTypeHint { path, data_type, - nullable: true, - default_constraint: None, inverted_index: false, }; assert!( @@ -955,8 +893,6 @@ mod tests { let type_hints = vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -989,8 +925,6 @@ mod tests { let type_hints = vec![JsonTypeHint { path: vec!["count".to_string()], data_type: ConcreteDataType::uint64_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1021,12 +955,10 @@ mod tests { } #[test] - fn test_encode_json_fills_missing_type_hint_with_default() -> Result<()> { + fn test_encode_json_fills_missing_type_hint_with_null() -> Result<()> { let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(7))), inverted_index: false, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1043,7 +975,7 @@ mod tests { let Value::Struct(user) = struct_field_value(&root, "user") else { panic!("Expected user Struct value"); }; - assert_eq!(struct_field_value(user, "age"), &Value::Int64(7)); + assert_eq!(struct_field_value(user, "age"), &Value::Null); Ok(()) } @@ -1052,8 +984,6 @@ mod tests { let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1075,21 +1005,15 @@ mod tests { } #[test] - fn test_encode_json_rejects_missing_non_null_type_hint() -> Result<()> { + fn test_encode_json_allows_missing_type_hint() -> Result<()> { let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }]; let settings = JsonSettings::try_new(type_hints, None)?; - let err = settings.encode(json!({})).unwrap_err(); - assert!( - err.to_string() - .contains("missing non-null JSON2 type hint path user.age") - ); + settings.encode(json!({}))?; Ok(()) } @@ -1099,17 +1023,11 @@ mod tests { JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(7))), inverted_index: false, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: false, - default_constraint: Some(ColumnDefaultConstraint::Value(Value::String( - "unknown".into(), - ))), inverted_index: false, }, ]; @@ -1123,11 +1041,8 @@ mod tests { let Value::Struct(user) = struct_field_value(&root, "user") else { panic!("Expected user Struct value"); }; - assert_eq!(struct_field_value(user, "age"), &Value::Int64(7)); - assert_eq!( - struct_field_value(user, "name"), - &Value::String("unknown".into()) - ); + assert_eq!(struct_field_value(user, "age"), &Value::Null); + assert_eq!(struct_field_value(user, "name"), &Value::Null); Ok(()) } diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index e8b75dee259..08a2fd86555 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -701,8 +701,6 @@ mod test { vec![JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(0), diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 3a9c1587966..2b9c44a326a 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -1060,22 +1060,16 @@ mod tests { JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, JsonTypeHint { path: vec!["commit".to_string(), "operation".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, JsonTypeHint { path: vec!["time_us".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ], @@ -1171,8 +1165,6 @@ mod tests { vec![JsonTypeHint { path: vec!["hint".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(2), @@ -1338,8 +1330,6 @@ mod tests { vec![JsonTypeHint { path: vec!["nested".to_string(), "hinted".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(8), diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index 8eec0e11520..f381f80722e 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -280,8 +280,6 @@ fn select_dynamic_hints( .map(|(path, stat)| JsonTypeHint { path: path.iter().map(|x| (*x).to_owned()).collect(), data_type: ConcreteDataType::from_arrow_type(&stat.data_type.as_arrow_type()), - nullable: true, - default_constraint: None, inverted_index: false, }) .collect() @@ -372,8 +370,6 @@ mod tests { vec![JsonTypeHint { path: vec!["hint".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(2), @@ -430,8 +426,6 @@ mod tests { vec![JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(0), @@ -486,8 +480,6 @@ mod tests { vec![JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(0), @@ -508,15 +500,11 @@ mod tests { JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, JsonTypeHint { path: vec!["source_only".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ], @@ -583,8 +571,6 @@ mod tests { vec![JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(0), @@ -594,15 +580,11 @@ mod tests { JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, JsonTypeHint { path: vec!["promoted".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ], diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 155791f8ab6..ea2a1225f13 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -2032,8 +2032,6 @@ mod tests { vec![JsonTypeHint { path: vec!["id".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], Some(0), diff --git a/src/pipeline/src/etl/transform.rs b/src/pipeline/src/etl/transform.rs index c9a1510735c..0869a4be0fc 100644 --- a/src/pipeline/src/etl/transform.rs +++ b/src/pipeline/src/etl/transform.rs @@ -22,8 +22,7 @@ use api::v1::ColumnDataType; use api::v1::value::ValueData; use chrono::Utc; use datatypes::json::{JsonSettings, JsonTypeHint}; -use datatypes::schema::{ColumnDefaultConstraint, FulltextOptions, SkippingIndexOptions}; -use datatypes::value::Value; +use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use snafu::{OptionExt, ResultExt, ensure}; use sql::parsers::utils::{ validate_column_fulltext_create_option, validate_column_skipping_index_create_option, @@ -56,7 +55,6 @@ const TRANSFORM_ON_FAILURE: &str = "on_failure"; const JSON2_TYPE: &str = "json2"; const JSON2_TYPE_HINT: &str = "type.json2[]"; const JSON2_TYPE_HINT_PATH: &str = "path"; -const JSON2_TYPE_HINT_NULLABLE: &str = "nullable"; pub use transformer::greptime::GreptimeTransformer; @@ -480,8 +478,6 @@ fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { })?; let mut path = None; let mut type_name = None; - let mut nullable = true; - let mut default = None; let mut index = None; for (key, value) in config { @@ -491,10 +487,6 @@ fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { match key { JSON2_TYPE_HINT_PATH => path = Some(yaml_string(value, JSON2_TYPE_HINT_PATH)?), TRANSFORM_TYPE => type_name = Some(yaml_string(value, TRANSFORM_TYPE)?), - JSON2_TYPE_HINT_NULLABLE => { - nullable = yaml_bool(value, JSON2_TYPE_HINT_NULLABLE)?; - } - TRANSFORM_DEFAULT => default = Some(value), TRANSFORM_INDEX => index = Some(value), _ => { return InvalidJson2TypeHintSnafu { @@ -528,13 +520,6 @@ fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { } ); let data_type = ColumnDataTypeWrapper::new(type_, None).into(); - let default_constraint = default - .map(|value| parse_json2_type_hint_default(value, &type_)) - .transpose()?; - if let Some(default_constraint) = &default_constraint { - default_constraint.validate(&data_type, nullable)?; - } - let inverted_index = if let Some(value) = index { let (index, options) = parse_transform_index(value)?; ensure!( @@ -552,40 +537,10 @@ fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { Ok(JsonTypeHint { path, data_type, - nullable, - default_constraint, inverted_index, }) } -fn parse_json2_type_hint_default( - value: &yaml_rust::Yaml, - type_: &ColumnDataType, -) -> Result { - if value.is_null() { - return Ok(ColumnDefaultConstraint::Value(Value::Null)); - } - - let value = match value { - yaml_rust::Yaml::Real(value) | yaml_rust::Yaml::String(value) => value.clone(), - yaml_rust::Yaml::Integer(value) => value.to_string(), - yaml_rust::Yaml::Boolean(value) => value.to_string(), - _ => { - return FieldMustBeTypeSnafu { - field: TRANSFORM_DEFAULT, - ty: "scalar", - } - .fail(); - } - }; - let value = api::v1::Value { - value_data: Some(parse_str_value(type_, &value)?), - }; - Ok(ColumnDefaultConstraint::Value( - api::helper::pb_value_to_value_ref(&value, None).into(), - )) -} - impl TryFrom<&yaml_rust::yaml::Hash> for Transform { type Error = Error; @@ -708,8 +663,6 @@ type: json2: - path: "user.id" type: int64 - nullable: false - default: 7 index: type: inverted - path: 'attrs."http.status_code"' @@ -726,19 +679,13 @@ type: hints[0].data_type, datatypes::prelude::ConcreteDataType::int64_datatype() ); - assert!(!hints[0].nullable); - assert_eq!( - hints[0].default_constraint, - Some(ColumnDefaultConstraint::Value(Value::Int64(7))) - ); assert!(hints[0].inverted_index); assert_eq!(hints[1].path, ["attrs", "http.status_code"]); - assert!(hints[1].nullable); } #[test] - fn test_transform_rejects_non_finite_json2_default() { - for default in ["NaN", "1e9999"] { + fn test_transform_rejects_json2_nullable_and_default() { + for option in ["nullable: false", "default: 7"] { let err = parse_transform(&format!( r#" field: payload @@ -746,22 +693,12 @@ type: json2: - path: score type: float64 - default: {default} + {option} "#, )) .unwrap_err(); - assert!( - matches!( - &err, - Error::Datatypes { - source: datatypes::error::Error::InvalidJson2Settings { .. }, - .. - } - ), - "{err:?}" - ); - assert!(err.to_string().contains("must be finite"), "{err}"); + assert!(err.to_string().contains("unsupported field"), "{err}"); } } diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index c85510e3e40..bbe3949a540 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -1101,8 +1101,6 @@ mod tests { vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }], None, @@ -1118,8 +1116,6 @@ mod tests { vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }], None, diff --git a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs index 9f01af708b2..a63b5ca076e 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs @@ -773,8 +773,6 @@ mod tests { vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), - nullable: false, - default_constraint: None, inverted_index: false, }], None, diff --git a/src/pipeline/tests/json_parse.rs b/src/pipeline/tests/json_parse.rs index b0c32950518..11c42f0fe4c 100644 --- a/src/pipeline/tests/json_parse.rs +++ b/src/pipeline/tests/json_parse.rs @@ -134,7 +134,6 @@ transform: json2: - path: "commitAuthor" type: string - nullable: false "#; let output = common::parse_and_exec(INPUT_VALUE_OBJ, pipeline_yaml); diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index d5a89f54076..ca38ec0d65e 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -1565,8 +1565,8 @@ ALTER TABLE metrics REPARTITION MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, - user.id STRING NOT NULL, - user.name STRING DEFAULT 'anonymous', + user.id STRING, + user.name STRING, request_id STRING INVERTED INDEX )"#; let mut statements = @@ -1590,8 +1590,6 @@ MODIFY COLUMN attrs JSON2 ( assert_eq!(Some(2000), options.max_auto_expanded_paths); assert_eq!(4, options.type_hints.len()); assert_eq!(vec!["user", "id"], options.type_hints[1].path); - assert!(!options.type_hints[1].nullable); - assert!(options.type_hints[2].default.is_some()); assert!(options.type_hints[3].inverted_index); let formatted = alter_table.to_string(); diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index fadc37c6d2f..2bf3c9407a6 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -15,7 +15,7 @@ use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; use datatypes::json::JSON2_MAX_STRUCTURED_DEPTH; use snafu::{ResultExt, ensure}; -use sqlparser::ast::{DataType, ExactNumberInfo, Expr, ObjectName, UnaryOperator}; +use sqlparser::ast::{DataType, ExactNumberInfo, ObjectName}; use sqlparser::dialect::keywords::Keyword; use sqlparser::parser::Parser; use sqlparser::tokenizer::Token; @@ -157,49 +157,22 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { let data_type = parser.parse_data_type().context(SyntaxSnafu)?; let data_type = normalize_json2_type_hint_type(data_type)?; - let mut nullable = true; - let mut nullable_set = false; - let mut default = None; let mut inverted_index = false; loop { - if parser.parse_keywords(&[Keyword::NOT, Keyword::NULL]) { - ensure!( - !nullable_set, - InvalidSqlSnafu { - msg: format!( - "NULL/NOT NULL option already specified for JSON2 type hint '{}'", - path.join(".") - ) - } - ); - nullable = false; - nullable_set = true; - } else if parser.parse_keyword(Keyword::NULL) { - ensure!( - !nullable_set, - InvalidSqlSnafu { - msg: format!( - "NULL/NOT NULL option already specified for JSON2 type hint '{}'", - path.join(".") - ) - } - ); - nullable = true; - nullable_set = true; + if parser.parse_keywords(&[Keyword::NOT, Keyword::NULL]) + || parser.parse_keyword(Keyword::NULL) + { + return InvalidSqlSnafu { + msg: "JSON2 type hint NULL/NOT NULL is not supported; hinted fields are always nullable" + .to_string(), + } + .fail(); } else if parser.parse_keyword(Keyword::DEFAULT) { - ensure!( - default.is_none(), - InvalidSqlSnafu { - msg: format!( - "duplicated DEFAULT option for JSON2 type hint '{}'", - path.join(".") - ) - } - ); - let expr = parser.parse_expr().context(SyntaxSnafu)?; - ensure_json2_default_expr_is_literal(&expr)?; - default = Some(expr); + return InvalidSqlSnafu { + msg: "JSON2 type hint DEFAULT is not supported".to_string(), + } + .fail(); } else if let Token::Word(word) = parser.peek_token().token && word.value.eq_ignore_ascii_case(INVERTED) { @@ -242,8 +215,6 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { Ok(JsonTypeHint { path, data_type, - nullable, - default, inverted_index, }) } @@ -298,24 +269,6 @@ fn normalize_json2_type_hint_type(data_type: DataType) -> Result { Ok(normalized) } -fn ensure_json2_default_expr_is_literal(expr: &Expr) -> Result<()> { - let is_literal = match expr { - Expr::Value(_) => true, - Expr::UnaryOp { op, expr } => { - matches!(op, UnaryOperator::Plus | UnaryOperator::Minus) - && matches!(expr.as_ref(), Expr::Value(_)) - } - _ => false, - }; - ensure!( - is_literal, - InvalidSqlSnafu { - msg: "JSON2 type hint DEFAULT only supports literal values", - } - ); - Ok(()) -} - fn ensure_no_path_conflict(hints: &[JsonTypeHint], path: &[String]) -> Result<()> { for hint in hints { ensure!( @@ -375,10 +328,10 @@ mod tests { r#" CREATE TABLE traces ( log_json_data JSON2 ( - "service.name" STRING NOT NULL DEFAULT 'null' INVERTED INDEX, - http.method STRING NOT NULL, - status_code INT64 NOT NULL, - comment STRING NULL, + "service.name" STRING INVERTED INDEX, + http.method STRING, + status_code INT64, + comment STRING, ), ts TIMESTAMP TIME INDEX, )"#, @@ -393,29 +346,17 @@ CREATE TABLE traces ( assert_eq!(hints[0].path, vec!["service.name"]); assert_eq!(hints[0].data_type, DataType::String(None)); - assert!(!hints[0].nullable); - assert_eq!( - hints[0] - .default - .as_ref() - .map(|expr| expr.to_string()) - .as_deref(), - Some("'null'") - ); assert!(hints[0].inverted_index); assert_eq!(hints[1].path, vec!["http", "method"]); assert_eq!(hints[1].data_type, DataType::String(None)); - assert!(!hints[1].nullable); assert!(!hints[1].inverted_index); assert_eq!(hints[2].path, vec!["status_code"]); assert_eq!(hints[2].data_type, DataType::BigInt(None)); - assert!(!hints[2].nullable); assert_eq!(hints[3].path, vec!["comment"]); assert_eq!(hints[3].data_type, DataType::String(None)); - assert!(hints[3].nullable); } #[test] @@ -484,7 +425,7 @@ CREATE TABLE traces ( } #[test] - fn test_parse_json2_type_hint_default_nullable() { + fn test_parse_json2_type_hint_defaults_to_nullable() { let column = parse_json2_column( r#" CREATE TABLE traces ( @@ -495,7 +436,7 @@ CREATE TABLE traces ( let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 1); - assert!(hints[0].nullable); + assert_eq!(hints[0].data_type, DataType::String(None)); } #[test] @@ -560,57 +501,19 @@ CREATE TABLE traces ( } #[test] - fn test_parse_json2_type_hint_default_accepts_signed_literals() { - let column = parse_json2_column( - r#" -CREATE TABLE traces ( - log_json_data JSON2 ( - negative_int INT64 DEFAULT -5, - positive_float FLOAT64 DEFAULT +1.5 - ), - ts TIMESTAMP TIME INDEX, -)"#, - ); - - let hints = column.extensions.json2_options.unwrap().type_hints; - assert_eq!(hints.len(), 2); - assert_eq!( - hints[0] - .default - .as_ref() - .map(|expr| expr.to_string()) - .as_deref(), - Some("-5") - ); - assert_eq!( - hints[1] - .default - .as_ref() - .map(|expr| expr.to_string()) - .as_deref(), - Some("+1.5") - ); - } - - #[test] - fn test_parse_json2_type_hint_default_rejects_function() { - let result = ParserContext::create_with_dialect( - r#" -CREATE TABLE traces ( - log_json_data JSON2 (status_code INT64 DEFAULT abs(-1)), - ts TIMESTAMP TIME INDEX, -)"#, - &GreptimeDbDialect {}, - ParseOptions::default(), - ); - - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("DEFAULT only supports literal values") - ); + fn test_parse_json2_type_hint_rejects_default() { + for default in ["-5", "abs(-1)"] { + let sql = format!( + "CREATE TABLE traces (log_json_data JSON2 (status_code INT64 DEFAULT {default}), ts TIMESTAMP TIME INDEX)" + ); + let err = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap_err(); + assert!(err.to_string().contains("DEFAULT is not supported")); + } } #[test] @@ -646,7 +549,7 @@ CREATE TABLE traces ( } #[test] - fn test_parse_json2_type_hint_rejects_duplicated_nullability() { + fn test_parse_json2_type_hint_rejects_nullability() { for sql in [ r#" CREATE TABLE traces ( @@ -680,7 +583,7 @@ CREATE TABLE traces ( result .unwrap_err() .to_string() - .contains("NULL/NOT NULL option already specified") + .contains("NULL/NOT NULL is not supported") ); } } diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index cf0b2486750..808e30f842c 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -678,8 +678,6 @@ mod tests { vec![datatypes::json::JsonTypeHint { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], None, diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index a4736d0a7aa..46b5b07c9f7 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -16,30 +16,27 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter}; use common_catalog::consts::FILE_ENGINE; -use common_sql::default_constraint::parse_column_default_constraint; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ - ColumnDefaultConstraint, FulltextOptions, SkippingIndexOptions, VectorDistanceMetric, - VectorIndexEngineType, VectorIndexOptions, + FulltextOptions, SkippingIndexOptions, VectorDistanceMetric, VectorIndexEngineType, + VectorIndexOptions, }; use itertools::Itertools; use serde::Serialize; use snafu::ResultExt; -use sqlparser::ast::{ColumnOption, ColumnOptionDef, DataType, Expr}; +use sqlparser::ast::{ColumnOptionDef, DataType, Expr}; use sqlparser_derive::{Visit, VisitMut}; use crate::ast::{ColumnDef, Ident, ObjectName, Value as SqlValue}; -use crate::dialect::GreptimeDbDialect; use crate::error::{ InvalidFlowQuerySnafu, InvalidSqlSnafu, Result, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, }; -use crate::parser::ParserContext; use crate::statements::query::Query as GtQuery; use crate::statements::statement::Statement; use crate::statements::tql::Tql; -use crate::statements::{OptionMap, sql_data_type_to_concrete_data_type, value_to_sql_value}; +use crate::statements::{OptionMap, sql_data_type_to_concrete_data_type}; const LINE_SEP: &str = ",\n"; const COMMA_SEP: &str = ", "; @@ -164,8 +161,6 @@ impl Display for Json2Options { pub struct JsonTypeHint { pub path: Vec, pub data_type: DataType, - pub nullable: bool, - pub default: Option, pub inverted_index: bool, } @@ -363,8 +358,6 @@ impl ColumnExtensions { Ok(datatypes::json::JsonTypeHint { path: hint.path.clone(), data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - nullable: hint.nullable, - default_constraint: build_json_type_hint_default_constraint(hint)?, inverted_index: hint.inverted_index, }) }) @@ -384,15 +377,9 @@ impl ColumnExtensions { .into_iter() .map(|hint| { let data_type = json_type_hint_sql_data_type(&hint.data_type)?; - let default = hint - .default_constraint - .map(|constraint| column_default_constraint_to_expr(&constraint)) - .transpose()?; Ok(JsonTypeHint { path: hint.path, data_type, - nullable: hint.nullable, - default, inverted_index: hint.inverted_index, }) }) @@ -406,38 +393,6 @@ impl ColumnExtensions { } } -fn build_json_type_hint_default_constraint( - hint: &JsonTypeHint, -) -> Result> { - let Some(default) = &hint.default else { - return Ok(None); - }; - - let data_type = json_type_hint_concrete_data_type(&hint.data_type)?; - let opts = [ColumnOptionDef { - name: None, - option: ColumnOption::Default(default.clone()), - }]; - - // Use the JSON path as the column name context for default value parsing errors. - let json_path = hint.path.join("."); - let default_constraint = parse_column_default_constraint(&json_path, &data_type, &opts, None) - .context(crate::error::SqlCommonSnafu)?; - - if let Some(constraint) = &default_constraint { - constraint - .validate(&data_type, hint.nullable) - .map_err(|e| { - InvalidSqlSnafu { - msg: format!("invalid DEFAULT for JSON2 type hint '{}': {e}", json_path), - } - .build() - })?; - } - - Ok(default_constraint) -} - fn json_type_hint_concrete_data_type(data_type: &DataType) -> Result { let data_type = sql_data_type_to_concrete_data_type(data_type)?; normalize_json_type_hint_concrete_data_type(&data_type) @@ -483,36 +438,18 @@ fn json_type_hint_sql_data_type(data_type: &ConcreteDataType) -> Result Result { - match constraint { - ColumnDefaultConstraint::Value(value) => Ok(Expr::Value(value_to_sql_value(value)?.into())), - ColumnDefaultConstraint::Function(function) => { - ParserContext::parse_function(function, &GreptimeDbDialect {}) - } - } -} - fn format_json_type_hint(hint: &JsonTypeHint) -> String { let path = hint .path .iter() .map(|segment| format_json_path_segment(segment)) .join("."); - let nullability = if hint.nullable { " NULL" } else { " NOT NULL" }; - let default = hint - .default - .as_ref() - .map(|expr| format!(" DEFAULT {expr}")) - .unwrap_or_default(); let inverted_index = if hint.inverted_index { " INVERTED INDEX" } else { "" }; - format!( - "{} {}{}{}{}", - path, hint.data_type, nullability, default, inverted_index - ) + format!("{} {}{}", path, hint.data_type, inverted_index) } fn format_json_path_segment(segment: &str) -> String { @@ -817,8 +754,6 @@ mod tests { use datatypes::json::{JsonSettings, JsonTypeHint as DatatypeJsonTypeHint}; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::ColumnDefaultConstraint; - use datatypes::value::Value; use super::*; use crate::dialect::GreptimeDbDialect; @@ -995,7 +930,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "service.name" STRING, - "a.b"."c" INT64 NOT NULL, + "a.b"."c" INT64, a."b.c" STRING ), ts TIMESTAMP TIME INDEX @@ -1011,9 +946,9 @@ ENGINE=mito r#" CREATE TABLE traces ( log_json_data JSON2( - "service.name" STRING NULL, - "a.b"."c" BIGINT NOT NULL, - "a"."b.c" STRING NULL + "service.name" STRING, + "a.b"."c" BIGINT, + "a"."b.c" STRING ), ts TIMESTAMP NOT NULL, TIME INDEX (ts) @@ -1039,7 +974,7 @@ ENGINE=mito fn test_parse_json2_max_auto_expanded_paths_option() -> Result<()> { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64 NOT NULL, + status_code INT64, max_auto_expanded_paths = 1 ), ts TIMESTAMP TIME INDEX @@ -1065,7 +1000,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "1abc" STRING, - a."2b" INT64 NOT NULL + a."2b" INT64 ), ts TIMESTAMP TIME INDEX )"#; @@ -1080,8 +1015,8 @@ ENGINE=mito r#" CREATE TABLE traces ( log_json_data JSON2( - "1abc" STRING NULL, - "a"."2b" BIGINT NOT NULL + "1abc" STRING, + "a"."2b" BIGINT ), ts TIMESTAMP NOT NULL, TIME INDEX (ts) @@ -1104,7 +1039,7 @@ ENGINE=mito } #[test] - fn test_json2_type_hint_default_builds_default_constraint() { + fn test_json2_type_hint_rejects_default() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( status_code INT64 DEFAULT -5, @@ -1114,67 +1049,24 @@ ENGINE=mito ), ts TIMESTAMP TIME INDEX )"#; - let result = + let err = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - let Statement::CreateTable(create_table) = &result[0] else { - unreachable!() - }; - let settings = create_table.columns[0] - .extensions - .build_json_settings() - .unwrap() - .unwrap(); - let hints = settings.type_hints(); - - assert_eq!(hints[0].data_type, ConcreteDataType::int64_datatype()); - assert_eq!( - hints[0].default_constraint, - Some(ColumnDefaultConstraint::Value(Value::Int64(-5))) - ); - assert_eq!(hints[1].data_type, ConcreteDataType::float64_datatype()); - assert_eq!( - hints[1].default_constraint, - Some(ColumnDefaultConstraint::Value(Value::Float64(1.5.into()))) - ); - assert_eq!(hints[2].data_type, ConcreteDataType::boolean_datatype()); - assert_eq!( - hints[2].default_constraint, - Some(ColumnDefaultConstraint::Value(Value::Boolean(false))) - ); - assert_eq!(hints[3].data_type, ConcreteDataType::string_datatype()); - assert_eq!( - hints[3].default_constraint, - Some(ColumnDefaultConstraint::Value(Value::String( - "unknown".into() - ))) - ); + .unwrap_err(); + assert!(err.to_string().contains("DEFAULT is not supported")); } #[test] - fn test_json2_type_hint_not_null_default_null_is_rejected() { + fn test_json2_type_hint_rejects_not_null() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( status_code INT64 NOT NULL DEFAULT NULL ), ts TIMESTAMP TIME INDEX )"#; - let result = + let err = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - let Statement::CreateTable(create_table) = &result[0] else { - unreachable!() - }; - let err = create_table.columns[0] - .extensions - .build_json_settings() - .unwrap_err(); - assert!( - err.to_string() - .contains("Default value should not be null for non null column") - ); + .unwrap_err(); + assert!(err.to_string().contains("NULL/NOT NULL is not supported")); } #[test] @@ -1185,36 +1077,26 @@ ENGINE=mito DatatypeJsonTypeHint { path: vec!["i".to_string()], data_type: ConcreteDataType::int32_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["f".to_string()], data_type: ConcreteDataType::float32_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["u".to_string()], data_type: ConcreteDataType::uint32_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["s".to_string()], data_type: ConcreteDataType::string_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["b".to_string()], data_type: ConcreteDataType::boolean_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }, ], @@ -1241,8 +1123,6 @@ ENGINE=mito vec![DatatypeJsonTypeHint { path: vec!["u".to_string()], data_type: ConcreteDataType::date_datatype(), - nullable: true, - default_constraint: None, inverted_index: false, }], None, diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index 493b1d85c16..ab529a20971 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -423,7 +423,7 @@ show create table json2_finite_paths; | | "ts" TIMESTAMP(3) NOT NULL, | | | "j" JSON2( | | | max_auto_expanded_paths = 1, | -| | "hint" STRING NULL | +| | "hint" STRING | | | ) NULL, | | | TIME INDEX ("ts") | | | ) | diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index 5662907d1d7..c23f9e2533f 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -27,8 +27,8 @@ Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, - user.id STRING NOT NULL, - user.name STRING DEFAULT 'anonymous', + user.id STRING, + user.name STRING, request_id STRING INVERTED INDEX ); @@ -38,8 +38,8 @@ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, - user.id STRING NOT NULL, - user.name STRING DEFAULT 'anonymous', + user.id STRING, + user.name STRING, request_id STRING INVERTED INDEX ); diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql index 2c17138b86d..c89556c3485 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -19,8 +19,8 @@ ALTER TABLE application_logs ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, - user.id STRING NOT NULL, - user.name STRING DEFAULT 'anonymous', + user.id STRING, + user.name STRING, request_id STRING INVERTED INDEX ); @@ -28,8 +28,8 @@ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, - user.id STRING NOT NULL, - user.name STRING DEFAULT 'anonymous', + user.id STRING, + user.name STRING, request_id STRING INVERTED INDEX ); diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.result b/tests/cases/standalone/common/types/json/json2_type_hints.result index 4e997cb1451..7739d1d1f8a 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.result +++ b/tests/cases/standalone/common/types/json/json2_type_hints.result @@ -1,10 +1,10 @@ CREATE TABLE json2_type_hints ( ts TIMESTAMP TIME INDEX, j JSON2 ( - user.age BIGINT NOT NULL DEFAULT 18, - user.name STRING DEFAULT 'unknown', - user.active BOOLEAN NULL, - score DOUBLE NULL DEFAULT 1.5 + user.age BIGINT, + user.name STRING, + user.active BOOLEAN, + score DOUBLE ) ) WITH ( 'append_mode' = 'true' @@ -14,26 +14,26 @@ Affected Rows: 0 SHOW CREATE TABLE json2_type_hints; -+------------------+--------------------------------------------------+ -| Table | Create Table | -+------------------+--------------------------------------------------+ -| json2_type_hints | CREATE TABLE IF NOT EXISTS "json2_type_hints" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "j" JSON2( | -| | max_auto_expanded_paths = 100, | -| | "user"."age" BIGINT NOT NULL DEFAULT 18, | -| | "user"."name" STRING NULL DEFAULT 'unknown', | -| | "user"."active" BOOLEAN NULL, | -| | "score" DOUBLE NULL DEFAULT 1.5 | -| | ) NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | append_mode = 'true' | -| | ) | -+------------------+--------------------------------------------------+ ++------------------+-------------------------------------------------+ +| Table | Create Table | ++------------------+-------------------------------------------------+ +| json2_type_hints | CREATE TABLE IF NOT EXISTS "json2_type_hints" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "user"."age" BIGINT, | +| | "user"."name" STRING, | +| | "user"."active" BOOLEAN, | +| | "score" DOUBLE | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++------------------+-------------------------------------------------+ INSERT INTO json2_type_hints VALUES @@ -55,8 +55,8 @@ ORDER BY ts; | json_get(json2_type_hints.j,Utf8("$.user.age")) | json_get(json2_type_hints.j,Utf8("$.user.name")) | json_get(json2_type_hints.j,Utf8("$.user.active")) | json_get(json2_type_hints.j,Utf8("$.score")) | +-------------------------------------------------+--------------------------------------------------+----------------------------------------------------+----------------------------------------------+ | 42 | Alice | true | 3.25 | -| 18 | Bob | | 1.5 | -| 18 | unknown | | 1.5 | +| | Bob | | | +| | | | | +-------------------------------------------------+--------------------------------------------------+----------------------------------------------------+----------------------------------------------+ INSERT INTO json2_type_hints @@ -64,22 +64,6 @@ VALUES (4, '{"user":{"age":"bad"}}'); Error: 1004(InvalidArguments), Invalid JSON: JSON value at user.age does not match JSON2 type hint Int64 -CREATE TABLE json2_type_hints_required ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - user.age BIGINT NOT NULL - ) -) WITH ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -INSERT INTO json2_type_hints_required -VALUES (1, '{}'); - -Error: 1004(InvalidArguments), Invalid JSON: missing non-null JSON2 type hint path user.age - CREATE TABLE json2_type_hints_timestamp ( ts TIMESTAMP TIME INDEX, j JSON2 ( @@ -89,30 +73,6 @@ CREATE TABLE json2_type_hints_timestamp ( Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP -CREATE TABLE json2_default_null_ok ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - a BIGINT NULL DEFAULT NULL - ) -) WITH ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -DROP TABLE json2_default_null_ok; - -Affected Rows: 0 - -CREATE TABLE json2_default_null_check ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - a BIGINT NOT NULL DEFAULT NULL - ) -); - -Error: 2000(InvalidSyntax), Invalid SQL, error: invalid DEFAULT for JSON2 type hint 'a': Default value should not be null for non null column - -- A type hint at the maximum supported depth is accepted. CREATE TABLE json2_type_hint_depth_50 ( ts TIMESTAMP TIME INDEX, @@ -141,10 +101,6 @@ DROP TABLE json2_type_hints; Affected Rows: 0 -DROP TABLE json2_type_hints_required; - -Affected Rows: 0 - DROP TABLE json2_type_hint_depth_50; Affected Rows: 0 diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.sql b/tests/cases/standalone/common/types/json/json2_type_hints.sql index 741684f79c5..aa88ae60bf0 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.sql +++ b/tests/cases/standalone/common/types/json/json2_type_hints.sql @@ -1,10 +1,10 @@ CREATE TABLE json2_type_hints ( ts TIMESTAMP TIME INDEX, j JSON2 ( - user.age BIGINT NOT NULL DEFAULT 18, - user.name STRING DEFAULT 'unknown', - user.active BOOLEAN NULL, - score DOUBLE NULL DEFAULT 1.5 + user.age BIGINT, + user.name STRING, + user.active BOOLEAN, + score DOUBLE ) ) WITH ( 'append_mode' = 'true' @@ -29,18 +29,6 @@ ORDER BY ts; INSERT INTO json2_type_hints VALUES (4, '{"user":{"age":"bad"}}'); -CREATE TABLE json2_type_hints_required ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - user.age BIGINT NOT NULL - ) -) WITH ( - 'append_mode' = 'true' -); - -INSERT INTO json2_type_hints_required -VALUES (1, '{}'); - CREATE TABLE json2_type_hints_timestamp ( ts TIMESTAMP TIME INDEX, j JSON2 ( @@ -48,24 +36,6 @@ CREATE TABLE json2_type_hints_timestamp ( ) ); -CREATE TABLE json2_default_null_ok ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - a BIGINT NULL DEFAULT NULL - ) -) WITH ( - 'append_mode' = 'true' -); - -DROP TABLE json2_default_null_ok; - -CREATE TABLE json2_default_null_check ( - ts TIMESTAMP TIME INDEX, - j JSON2 ( - a BIGINT NOT NULL DEFAULT NULL - ) -); - -- A type hint at the maximum supported depth is accepted. CREATE TABLE json2_type_hint_depth_50 ( ts TIMESTAMP TIME INDEX, @@ -88,6 +58,4 @@ CREATE TABLE json2_type_hint_depth_51 ( DROP TABLE json2_type_hints; -DROP TABLE json2_type_hints_required; - DROP TABLE json2_type_hint_depth_50; From 9c2933de0b45f02011ecf4c144f7fb12169cd5cf Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 18 Sep 2026 07:28:40 +0000 Subject: [PATCH 04/63] feat(log-store): add the object store WAL batch, catalog and I/O (#9216) * feat(log-store): add the object store WAL batch, catalog and I/O Add the three modules between the object format and the store of the object store WAL: - batch: the open batch that accumulates admitted entries into the next object and assigns entry ids at admission. Ids are object-sequence-major, `(object_seq << 20) | position`, with positions starting at one per region per object, so a batch that is rolled back and admitted again under the same sequence hands out the same ids. The time the first entry was admitted is kept for the store's age-based sealing; an empty admission does not start it. - catalog: the in-memory index over object footers, by sequence and per region. Insertion is atomic and rejects an empty footer, duplicate region segments, invalid entry ranges, an already indexed sequence and entry ranges that are not strictly increasing across objects. The next object sequence continues after the largest indexed one and is raised above the largest entry id of any region, and is rejected once it no longer fits an entry id. - io: object store access under `/objects/`: a conditional create whose retry with identical content is a no-op and whose conflicting content is rejected, whole and range reads, and a listing of well-formed keys in sequence order. A prefix with a `.` or `..` component is rejected. When the store reports that the object exists but the read that compares its content fails, the read failure is returned with its retry hint instead of the create failure. The modules have no callers until the store lands, so they are declared with `#[allow(dead_code)]`. Only `entry_id` is exported. Signed-off-by: jeremyhi * test(log-store): drop comments that restate the batch test assertions Signed-off-by: jeremyhi * refactor(log-store): use pub(crate) in the object store WAL batch, catalog and I/O Signed-off-by: jeremyhi * fix(log-store): keep entry_id crate-private and narrow the create collision check Re-export `entry_id` as `pub(crate)`: nothing outside `log-store` uses it yet, and exporting it would freeze the raw `(object_seq, position)` encoding before the store owns id allocation. Treat only `ConditionNotMatch` as the sign that a conditional create collided with an existing object. A store may report `AlreadyExists` for an unrelated path, for example when a parent of the object is a file, and that failure must come back as the write failure rather than as the error of the read that would compare content. Signed-off-by: jeremyhi * refactor(log-store): rename the catalog's out_of_order helper to out_of_order_reason Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/Cargo.toml | 1 + src/log-store/src/error.rs | 79 ++- src/log-store/src/object_store_wal.rs | 24 +- src/log-store/src/object_store_wal/batch.rs | 359 +++++++++++ src/log-store/src/object_store_wal/catalog.rs | 468 +++++++++++++++ src/log-store/src/object_store_wal/io.rs | 558 ++++++++++++++++++ 6 files changed, 1480 insertions(+), 9 deletions(-) create mode 100644 src/log-store/src/object_store_wal/batch.rs create mode 100644 src/log-store/src/object_store_wal/catalog.rs create mode 100644 src/log-store/src/object_store_wal/io.rs diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index ee2c6a12dc0..15560b06b2c 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -50,6 +50,7 @@ common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-wal = { workspace = true, features = ["testing"] } itertools.workspace = true +object-store = { workspace = true, features = ["testing"] } rand.workspace = true rskafka = { workspace = true, features = ["unstable-fuzzing"] } uuid.workspace = true diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index dec7b42c3a5..54fcc84b690 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -315,6 +315,61 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display("Invalid WAL object store, {}", reason))] + InvalidWalObjectStore { + reason: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Invalid WAL entry range, region: {}, start: {}, end: {}", + region_id, + start_entry_id, + end_entry_id + ))] + InvalidWalEntryRange { + region_id: RegionId, + start_entry_id: u64, + end_entry_id: u64, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("WAL object sequence is exhausted, last sequence: {}", last_object_seq))] + WalObjectSequenceExhausted { + last_object_seq: u64, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "WAL entry positions of region {} in one object are exhausted", + region_id + ))] + WalEntryPositionExhausted { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("WAL object already exists with different content, path: {}", path))] + WalObjectConflict { + path: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to {} WAL object, path: {}", operation, path))] + WalObjectStore { + operation: &'static str, + path: String, + #[snafu(source)] + error: object_store::Error, + #[snafu(implicit)] + location: Location, + }, } pub type Result = std::result::Result; @@ -345,7 +400,9 @@ impl ErrorExt for Error { | IllegalNamespace { .. } | MissingKey { .. } | MissingValue { .. } - | OverrideCompactedEntry { .. } => StatusCode::InvalidArguments, + | OverrideCompactedEntry { .. } + | InvalidWalObjectStore { .. } + | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, StartWalTask { .. } | StopWalTask { .. } | IllegalState { .. } @@ -360,12 +417,17 @@ impl ErrorExt for Error { | WaitDumpIndex { .. } | MetaLengthExceededLimit { .. } => StatusCode::Internal, - CorruptedWalObject { .. } => StatusCode::Unexpected, + CorruptedWalObject { .. } + | WalObjectConflict { .. } + | WalObjectSequenceExhausted { .. } + | WalEntryPositionExhausted { .. } => StatusCode::Unexpected, // Object store related errors - CreateWriter { .. } | WriteIndex { .. } | ReadIndex { .. } | Io { .. } => { - StatusCode::StorageUnavailable - } + CreateWriter { .. } + | WriteIndex { .. } + | ReadIndex { .. } + | WalObjectStore { .. } + | Io { .. } => StatusCode::StorageUnavailable, // Raft engine FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => { StatusCode::StorageUnavailable @@ -391,9 +453,10 @@ impl ErrorExt for Error { use Error::*; match self { - CreateWriter { error, .. } | WriteIndex { error, .. } | ReadIndex { error, .. } => { - retry_hint_from_opendal_error(error) - } + CreateWriter { error, .. } + | WriteIndex { error, .. } + | ReadIndex { error, .. } + | WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), Io { error, .. } => retry_hint_from_io_error(error), FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index f4cd7e081c7..98895869b81 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -29,7 +29,29 @@ //! range, byte range and CRC32. The fixed-size trailer points at the footer and //! carries the CRC32 of the footer and of the whole object, so a reader locates //! the footer by reading the fixed-length trailer at the end of the object. +//! +//! Object sequences increase monotonically within one prefix and may leave +//! gaps, so recovery continues after the largest sequence it indexed. An object +//! is created conditionally: rewriting a sequence with the content it already +//! holds is a no-op at the object store, while different content under a taken +//! sequence is a conflict. Recovery lists the objects, reads and verifies only +//! the header, trailer and footer of each, and indexes the footers in sequence +//! order to rebuild the object catalog, which rejects a sequence it already +//! holds. Segments are read and checksummed only when a read decodes them. +//! +//! Entry ids are object-sequence-major, see [`entry_id`]: the high bits of an +//! id name the object that holds the entry, the low bits its position among +//! the entries of its region in that object. -// The format has no callers until the store that writes and reads objects lands. +// These modules have no callers until the store that writes and reads objects lands. +#[allow(dead_code)] +mod batch; +#[allow(dead_code)] +mod catalog; #[allow(dead_code)] mod format; +#[allow(dead_code)] +mod io; + +#[allow(unused_imports)] +pub(crate) use batch::entry_id; diff --git a/src/log-store/src/object_store_wal/batch.rs b/src/log-store/src/object_store_wal/batch.rs new file mode 100644 index 00000000000..b07ac8e8cb0 --- /dev/null +++ b/src/log-store/src/object_store_wal/batch.rs @@ -0,0 +1,359 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Accumulation of admitted entries into the batch that becomes the next +//! object, and the entry id scheme. + +use std::collections::HashMap; +use std::time::Instant; + +use snafu::ensure; +use store_api::logstore::EntryId; +use store_api::logstore::entry::Entry; +use store_api::storage::RegionId; + +use crate::error::{Result, WalEntryPositionExhaustedSnafu}; + +/// Bits of an entry id that hold the position of the entry among the entries +/// of its region inside the object; the remaining high bits hold the object +/// sequence. +pub(crate) const POSITION_BITS: u32 = 20; +/// Positions are `1..POSITION_LIMIT`, so an object holds at most 2^20 - 1 +/// entries of one region and a batch seals before a region reaches the limit. +/// This is a theoretical bound: an object is sealed by size long before. +pub(crate) const POSITION_LIMIT: u64 = 1 << POSITION_BITS; +/// Object sequences are `0..OBJECT_SEQ_LIMIT`, the 44 bits an entry id leaves +/// above the position. +pub(crate) const OBJECT_SEQ_LIMIT: u64 = 1 << (u64::BITS - POSITION_BITS); + +/// Returns the id of the entry at `position` among the entries of its region +/// in the object `object_seq`. +/// +/// Ids are object-sequence-major: `entry_id >> 20` names the object that holds +/// the entry, and the low bits are its position in that object, which starts +/// at one so that id zero, the watermark of a region without entries, is never +/// assigned. A region's ids increase with the object sequence and have gaps +/// wherever other regions or other positions took the sequence. +pub(crate) fn entry_id(object_seq: u64, position: u64) -> EntryId { + debug_assert!(object_seq < OBJECT_SEQ_LIMIT && (1..POSITION_LIMIT).contains(&position)); + (object_seq << POSITION_BITS) | position +} + +/// Returns the smallest object sequence whose entry ids are all greater than +/// `entry_id`. Zero names no entry, so it needs no floor. +/// +/// An id assigned under the earlier contiguous scheme carries no object +/// information, but the floor keeps every new id above it just the same. +pub(crate) fn sequence_floor(entry_id: EntryId) -> u64 { + if entry_id == 0 { + 0 + } else { + (entry_id >> POSITION_BITS) + 1 + } +} + +/// Entries admitted since the last seal, together with the position of the +/// last entry admitted per region. +/// +/// Every entry is assigned its id at admission from the sequence the batch +/// takes when it is sealed, so a batch that is rolled back and admitted again +/// under the same sequence hands out the same ids. +#[derive(Debug)] +pub(crate) struct OpenBatch { + max_bytes: usize, + entries: Vec, + estimated_bytes: usize, + /// When the first entry of the batch was admitted. + first_admitted_at: Option, + positions: HashMap, +} + +impl OpenBatch { + pub(crate) fn new(max_bytes: usize) -> Self { + Self { + max_bytes, + entries: Vec::new(), + estimated_bytes: 0, + first_admitted_at: None, + positions: HashMap::new(), + } + } + + /// Returns true when admitting `entries` would take a region past the + /// position range, so that `entries` need a batch of their own. + pub(crate) fn would_exhaust_positions(&self, entries: &[Entry]) -> bool { + self.check_positions(entries).is_err() + } + + /// Admits `entries` into the batch that becomes the object `object_seq`, + /// assigning each the next position of its region, and returns the last id + /// assigned to every region in `entries`. Nothing is admitted when a + /// region would run past the position range. + pub(crate) fn admit( + &mut self, + object_seq: u64, + mut entries: Vec, + ) -> Result> { + self.check_positions(&entries)?; + let mut last_entry_ids = HashMap::new(); + for entry in &mut entries { + let region_id = entry.region_id(); + let position = self.positions.entry(region_id).or_insert(0); + *position += 1; + let entry_id = entry_id(object_seq, *position); + entry.set_entry_id(entry_id); + last_entry_ids.insert(region_id, entry_id); + } + self.estimated_bytes += entries.iter().map(Entry::estimated_size).sum::(); + if !entries.is_empty() { + self.first_admitted_at.get_or_insert_with(Instant::now); + } + self.entries.extend(entries); + Ok(last_entry_ids) + } + + /// Checks that every region in `entries` stays inside the position range + /// once they are admitted. + fn check_positions(&self, entries: &[Entry]) -> Result<()> { + let mut positions = HashMap::new(); + for entry in entries { + let region_id = entry.region_id(); + let position = positions + .entry(region_id) + .or_insert_with(|| self.positions.get(®ion_id).copied().unwrap_or(0)); + *position += 1; + ensure!( + *position < POSITION_LIMIT, + WalEntryPositionExhaustedSnafu { region_id } + ); + } + Ok(()) + } + + pub(crate) fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Returns the estimated size of the admitted entries. + pub(crate) fn estimated_bytes(&self) -> usize { + self.estimated_bytes + } + + /// Returns when the first admitted entry was admitted, if any. + pub(crate) fn first_admitted_at(&self) -> Option { + self.first_admitted_at + } + + /// Returns true once the admitted entries reach the size limit. + pub(crate) fn should_seal(&self) -> bool { + !self.is_empty() && self.estimated_bytes >= self.max_bytes + } + + /// Takes the admitted entries out of the batch together with the time the + /// first of them was admitted. The next admission starts at position one + /// again, under the next sequence. + pub(crate) fn seal(&mut self) -> (Vec, Instant) { + self.estimated_bytes = 0; + self.positions.clear(); + let first_admitted_at = self.first_admitted_at.take().unwrap_or_else(Instant::now); + (std::mem::take(&mut self.entries), first_admitted_at) + } + + /// Drops the admitted entries, so the next admission under the same + /// sequence hands out the same ids again. + pub(crate) fn reset(&mut self) { + let _ = self.seal(); + } +} + +#[cfg(test)] +mod tests { + use store_api::logstore::entry::NaiveEntry; + use store_api::logstore::provider::Provider; + + use super::*; + use crate::error::Error; + + fn entry(region_id: RegionId, payload_len: usize) -> Entry { + Entry::Naive(NaiveEntry { + provider: Provider::object_store_provider(region_id, "wal".to_string()), + region_id, + entry_id: 0, + data: vec![0; payload_len], + }) + } + + fn entry_ids(entries: &[Entry]) -> Vec<(RegionId, EntryId)> { + entries + .iter() + .map(|entry| (entry.region_id(), entry.entry_id())) + .collect() + } + + #[test] + fn test_entry_id_scheme() { + assert_eq!(1, entry_id(0, 1)); + assert_eq!(0x10_0001, entry_id(1, 1)); + // 7 << 52 | 3 + assert_eq!(31_525_197_391_593_475, entry_id(0x7_0000_0000, 3)); + assert_eq!(u64::MAX, entry_id(OBJECT_SEQ_LIMIT - 1, POSITION_LIMIT - 1)); + assert_eq!(1 << 44, OBJECT_SEQ_LIMIT); + + assert_eq!(0, sequence_floor(0)); + assert_eq!(1, sequence_floor(1)); + assert_eq!(1, sequence_floor(POSITION_LIMIT - 1)); + assert_eq!(2, sequence_floor(entry_id(1, 1))); + assert_eq!(8, sequence_floor(entry_id(7, POSITION_LIMIT - 1))); + // A contiguous id well past the first object is placed by its high bits. + assert_eq!(5, sequence_floor(5_000_000)); + assert_eq!(OBJECT_SEQ_LIMIT, sequence_floor(u64::MAX)); + } + + #[test] + fn test_batch_assigns_positions_per_region_under_the_object_sequence() { + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + let mut batch = OpenBatch::new(usize::MAX); + + let first = batch + .admit(5, vec![entry(region_a, 1), entry(region_b, 1)]) + .unwrap(); + assert_eq!( + HashMap::from([(region_a, entry_id(5, 1)), (region_b, entry_id(5, 1))]), + first + ); + let second = batch + .admit( + 5, + vec![entry(region_b, 1), entry(region_a, 1), entry(region_b, 1)], + ) + .unwrap(); + assert_eq!( + HashMap::from([(region_a, entry_id(5, 2)), (region_b, entry_id(5, 3))]), + second + ); + + assert_eq!( + vec![ + (region_a, entry_id(5, 1)), + (region_b, entry_id(5, 1)), + (region_b, entry_id(5, 2)), + (region_a, entry_id(5, 2)), + (region_b, entry_id(5, 3)), + ], + entry_ids(&batch.seal().0) + ); + assert!(batch.is_empty()); + assert_eq!( + HashMap::from([(region_a, entry_id(6, 1))]), + batch.admit(6, vec![entry(region_a, 1)]).unwrap() + ); + assert_eq!(POSITION_LIMIT, entry_id(6, 1) - entry_id(5, 1)); + } + + #[test] + fn test_batch_seals_at_size_limit() { + let region_id = RegionId::new(1, 1); + let first = entry(region_id, 8); + let second = entry(region_id, 8); + let max_bytes = first.estimated_size() + second.estimated_size(); + let mut batch = OpenBatch::new(max_bytes); + + assert!(!batch.should_seal()); + batch.admit(0, vec![first]).unwrap(); + assert!(!batch.should_seal()); + batch.admit(0, vec![second]).unwrap(); + assert!(batch.should_seal()); + + assert_eq!(2, batch.seal().0.len()); + assert!(!batch.should_seal()); + } + + #[test] + fn test_batch_admission_clock_starts_with_the_first_entry() { + let region_id = RegionId::new(1, 1); + let mut batch = OpenBatch::new(usize::MAX); + + assert!(batch.admit(0, Vec::new()).unwrap().is_empty()); + assert!(batch.is_empty()); + assert_eq!(None, batch.first_admitted_at()); + + let before = Instant::now(); + batch.admit(0, vec![entry(region_id, 1)]).unwrap(); + let first_admitted_at = batch.first_admitted_at().unwrap(); + assert!(first_admitted_at >= before); + batch.admit(0, vec![entry(region_id, 1)]).unwrap(); + assert_eq!(Some(first_admitted_at), batch.first_admitted_at()); + assert_eq!(first_admitted_at, batch.seal().1); + assert_eq!(None, batch.first_admitted_at()); + } + + #[test] + fn test_batch_reset_hands_out_the_same_ids_again() { + let region_id = RegionId::new(1, 1); + let mut batch = OpenBatch::new(usize::MAX); + + assert_eq!( + HashMap::from([(region_id, entry_id(3, 1))]), + batch.admit(3, vec![entry(region_id, 1)]).unwrap() + ); + batch.reset(); + assert!(batch.is_empty()); + assert_eq!( + HashMap::from([(region_id, entry_id(3, 1))]), + batch.admit(3, vec![entry(region_id, 1)]).unwrap() + ); + } + + #[test] + fn test_batch_refuses_a_region_past_the_position_range() { + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + let mut batch = OpenBatch::new(usize::MAX); + let entries = + |region_id, count: u64| (0..count).map(|_| entry(region_id, 0)).collect::>(); + + // An append that alone runs past the range fits no object. + let error = batch + .admit(0, entries(region_a, POSITION_LIMIT)) + .unwrap_err(); + assert!( + matches!(error, Error::WalEntryPositionExhausted { region_id, .. } if region_id == region_a), + "unexpected error: {error:?}" + ); + assert!(batch.is_empty()); + + // The range holds one entry fewer; the next entry of that region needs + // a new batch, while another region still fits. + let last = batch + .admit(0, entries(region_a, POSITION_LIMIT - 1)) + .unwrap(); + assert_eq!( + HashMap::from([(region_a, entry_id(0, POSITION_LIMIT - 1))]), + last + ); + assert!(batch.would_exhaust_positions(&entries(region_a, 1))); + assert!(!batch.would_exhaust_positions(&entries(region_b, 1))); + let error = batch.admit(0, entries(region_a, 1)).unwrap_err(); + assert!( + matches!(error, Error::WalEntryPositionExhausted { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(POSITION_LIMIT as usize - 1, batch.seal().0.len()); + assert_eq!( + HashMap::from([(region_a, entry_id(1, 1))]), + batch.admit(1, entries(region_a, 1)).unwrap() + ); + } +} diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs new file mode 100644 index 00000000000..5e9650d14a5 --- /dev/null +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -0,0 +1,468 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! In-memory index over the footers of the objects of one WAL prefix. + +use std::collections::BTreeMap; +use std::ops::Bound::{Excluded, Unbounded}; + +use snafu::{OptionExt, ensure}; +use store_api::storage::RegionId; + +use crate::error::{ + CorruptedWalObjectSnafu, InvalidWalEntryRangeSnafu, Result, WalObjectSequenceExhaustedSnafu, +}; +use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, sequence_floor}; +use crate::object_store_wal::format::FooterEntry; + +/// Indexes objects by sequence and, per region, the objects that hold entries +/// of that region. +#[derive(Debug, Default)] +pub(crate) struct ObjectCatalog { + objects: BTreeMap>, + regions: BTreeMap>, +} + +impl ObjectCatalog { + /// Indexes the footer of the object `object_seq`. Objects may be inserted + /// in any order, which lets recovery index them as it discovers them. + /// Inserting a sequence that is already indexed is rejected, whether or not + /// the footer matches the indexed one. + pub(crate) fn insert_object( + &mut self, + object_seq: u64, + mut footer: Vec, + ) -> Result<()> { + ensure!( + !footer.is_empty(), + CorruptedWalObjectSnafu { + reason: format!("object {object_seq} has an empty footer"), + } + ); + + footer.sort_unstable_by_key(|entry| entry.region_id); + for entries in footer.windows(2) { + ensure!( + entries[0].region_id != entries[1].region_id, + CorruptedWalObjectSnafu { + reason: format!( + "object {object_seq} has duplicate footer entries for region {}", + entries[0].region_id + ), + } + ); + } + for entry in &footer { + ensure!( + entry.entry_count > 0 && entry.min_entry_id <= entry.max_entry_id, + CorruptedWalObjectSnafu { + reason: format!( + "object {object_seq} has invalid entry range {}..={} with {} entries for region {}", + entry.min_entry_id, entry.max_entry_id, entry.entry_count, entry.region_id + ), + } + ); + } + + // Object sequences are unique: recovery indexes every listed key once, and + // accepting a repeated insertion would hide a caller that lost track of it. + // Retrying an identical write stays an object store concern. + ensure!( + !self.objects.contains_key(&object_seq), + CorruptedWalObjectSnafu { + reason: format!("object {object_seq} is already indexed"), + } + ); + + // Validate every region before mutating either index so insertion is atomic. + for entry in &footer { + let Some(region_objects) = self.regions.get(&entry.region_id) else { + continue; + }; + if let Some((&previous_seq, previous)) = region_objects.range(..object_seq).next_back() + { + ensure!( + previous.max_entry_id < entry.min_entry_id, + CorruptedWalObjectSnafu { + reason: out_of_order_reason( + entry.region_id, + previous_seq, + previous.max_entry_id, + object_seq, + entry.min_entry_id + ), + } + ); + } + if let Some((&next_seq, next)) = region_objects + .range((Excluded(object_seq), Unbounded)) + .next() + { + ensure!( + entry.max_entry_id < next.min_entry_id, + CorruptedWalObjectSnafu { + reason: out_of_order_reason( + entry.region_id, + object_seq, + entry.max_entry_id, + next_seq, + next.min_entry_id + ), + } + ); + } + } + + for entry in &footer { + self.regions + .entry(entry.region_id) + .or_default() + .insert(object_seq, entry.clone()); + } + self.objects.insert(object_seq, footer); + Ok(()) + } + + /// Returns the objects that hold entries of `region_id` overlapping + /// `start_entry_id..=end_entry_id`, ordered by object sequence. + pub(crate) fn objects_for_entry_range( + &self, + region_id: RegionId, + start_entry_id: u64, + end_entry_id: u64, + ) -> Result> { + ensure!( + start_entry_id <= end_entry_id, + InvalidWalEntryRangeSnafu { + region_id, + start_entry_id, + end_entry_id, + } + ); + + let Some(objects) = self.regions.get(®ion_id) else { + return Ok(Vec::new()); + }; + Ok(objects + .iter() + .filter(|(_, entry)| { + entry.max_entry_id >= start_entry_id && entry.min_entry_id <= end_entry_id + }) + .map(|(&object_seq, entry)| (object_seq, entry)) + .collect()) + } + + /// Returns the largest entry id indexed for `region_id`. + pub(crate) fn region_max_entry_id(&self, region_id: RegionId) -> Option { + self.regions + .get(®ion_id)? + .last_key_value() + .map(|(_, entry)| entry.max_entry_id) + } + + /// Returns the sequence to assign to the next object written after recovery. + /// + /// An empty catalog starts at zero, so the first object of a prefix always + /// takes sequence zero. Otherwise the sequence continues after the largest + /// indexed one, which recovery discovers regardless of insertion order, and + /// is raised further when the largest entry id of a region lies at or above + /// the ids that sequence would assign: ids assigned under the earlier + /// contiguous scheme carry no object information, and every new id of a + /// region must be greater than every id it already has. A sequence at or + /// above [`OBJECT_SEQ_LIMIT`] does not fit an entry id and is rejected. + pub(crate) fn next_object_seq(&self) -> Result { + let after_last = match self.objects.last_key_value() { + None => 0, + Some((&last_object_seq, _)) => last_object_seq + .checked_add(1) + .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, + }; + let floor = self + .regions + .keys() + .filter_map(|region_id| self.region_max_entry_id(*region_id)) + .map(sequence_floor) + .max() + .unwrap_or(0); + let next_object_seq = after_last.max(floor); + ensure!( + next_object_seq < OBJECT_SEQ_LIMIT, + WalObjectSequenceExhaustedSnafu { + last_object_seq: next_object_seq - 1, + } + ); + Ok(next_object_seq) + } + + /// Iterates over the indexed objects ordered by object sequence. + pub(crate) fn objects_in_order(&self) -> impl Iterator + '_ { + self.objects + .iter() + .map(|(&object_seq, footer)| (object_seq, footer.as_slice())) + } +} + +fn out_of_order_reason( + region_id: RegionId, + lower_object_seq: u64, + lower_max_entry_id: u64, + upper_object_seq: u64, + upper_min_entry_id: u64, +) -> String { + format!( + "entry ranges of region {region_id} are not strictly increasing, object {lower_object_seq} ends at {lower_max_entry_id}, object {upper_object_seq} starts at {upper_min_entry_id}" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error; + use crate::object_store_wal::batch::entry_id; + + #[test] + fn test_catalog_indexes_objects_and_queries_ranges() { + let region_one = RegionId::new(1, 1); + let region_two = RegionId::new(2, 1); + let mut catalog = ObjectCatalog::default(); + + // Recovery may discover objects out of order. + catalog + .insert_object( + 2, + vec![ + footer_entry(region_one, 4, 6), + footer_entry(region_two, 8, 9), + ], + ) + .unwrap(); + catalog + .insert_object(1, vec![footer_entry(region_one, 1, 3)]) + .unwrap(); + catalog + .insert_object(4, vec![footer_entry(region_one, 10, 12)]) + .unwrap(); + + assert_eq!(Some(12), catalog.region_max_entry_id(region_one)); + assert_eq!(Some(9), catalog.region_max_entry_id(region_two)); + assert_eq!(None, catalog.region_max_entry_id(RegionId::new(3, 1))); + + let objects = catalog.objects_for_entry_range(region_one, 3, 10).unwrap(); + assert_eq!(vec![1, 2, 4], object_seqs(&objects)); + assert_eq!( + vec![1, 2, 4], + catalog + .objects_in_order() + .map(|(object_seq, _)| object_seq) + .collect::>() + ); + } + + #[test] + fn test_catalog_rejects_duplicate_object_sequences() { + let region_id = RegionId::new(1, 1); + let mut catalog = ObjectCatalog::default(); + let first = footer_entry(region_id, 1, 2); + let second = footer_entry(RegionId::new(2, 1), 4, 5); + + catalog + .insert_object(1, vec![first.clone(), second.clone()]) + .unwrap(); + + // An identical footer is rejected just like a conflicting one. + assert_corrupted( + catalog.insert_object(1, vec![second, first.clone()]), + "object 1 is already indexed", + ); + + let mut conflicting = first; + conflicting.segment_offset += 1; + assert_corrupted( + catalog.insert_object(1, vec![conflicting]), + "object 1 is already indexed", + ); + assert_eq!(1, catalog.objects_in_order().count()); + } + + #[test] + fn test_catalog_resumes_object_sequence_after_recovery() { + let region_id = RegionId::new(1, 1); + let mut catalog = ObjectCatalog::default(); + assert_eq!(0, catalog.next_object_seq().unwrap()); + + // Recovery may discover objects out of order. + catalog + .insert_object(4, vec![footer_entry(region_id, 10, 12)]) + .unwrap(); + catalog + .insert_object(1, vec![footer_entry(region_id, 1, 3)]) + .unwrap(); + + assert_eq!(5, catalog.next_object_seq().unwrap()); + } + + #[test] + fn test_catalog_raises_object_sequence_above_existing_entry_ids() { + let region_one = RegionId::new(1, 1); + let region_two = RegionId::new(1, 2); + let mut catalog = ObjectCatalog::default(); + + // Ids that fit below the ids of the next sequence leave it alone. + catalog + .insert_object(0, vec![footer_entry(region_one, 1, 3)]) + .unwrap(); + catalog + .insert_object( + 1, + vec![footer_entry(region_two, entry_id(1, 1), entry_id(1, 2))], + ) + .unwrap(); + assert_eq!(2, catalog.next_object_seq().unwrap()); + + // A contiguous id past them names a later object: the sequence + // resumes above it, whichever region holds it. + catalog + .insert_object(2, vec![footer_entry(region_one, 5_000_000, 5_000_000)]) + .unwrap(); + assert_eq!(5, catalog.next_object_seq().unwrap()); + catalog + .insert_object( + 3, + vec![footer_entry(region_two, entry_id(7, 4), entry_id(7, 4))], + ) + .unwrap(); + assert_eq!(8, catalog.next_object_seq().unwrap()); + } + + #[test] + fn test_catalog_rejects_exhausted_object_sequence() { + let region_id = RegionId::new(1, 1); + let assert_exhausted = |catalog: &ObjectCatalog| { + let error = catalog.next_object_seq().unwrap_err(); + assert!( + error.to_string().contains("object sequence is exhausted"), + "unexpected error: {error}" + ); + }; + + // The last sequence that fits an entry id is indexed. + let mut catalog = ObjectCatalog::default(); + catalog + .insert_object(OBJECT_SEQ_LIMIT - 2, vec![footer_entry(region_id, 1, 2)]) + .unwrap(); + assert_eq!(OBJECT_SEQ_LIMIT - 1, catalog.next_object_seq().unwrap()); + catalog + .insert_object(OBJECT_SEQ_LIMIT - 1, vec![footer_entry(region_id, 3, 4)]) + .unwrap(); + assert_exhausted(&catalog); + + // A sequence that does not fit was written by an earlier scheme. + let mut catalog = ObjectCatalog::default(); + catalog + .insert_object(u64::MAX, vec![footer_entry(region_id, 1, 2)]) + .unwrap(); + assert_exhausted(&catalog); + + // An entry id that leaves no sequence above it. + let mut catalog = ObjectCatalog::default(); + catalog + .insert_object(0, vec![footer_entry(region_id, u64::MAX, u64::MAX)]) + .unwrap(); + assert_exhausted(&catalog); + } + + #[test] + fn test_catalog_rejects_overlapping_or_reversed_region_ranges() { + let region_id = RegionId::new(1, 1); + let mut catalog = ObjectCatalog::default(); + catalog + .insert_object(2, vec![footer_entry(region_id, 10, 20)]) + .unwrap(); + + assert_corrupted( + catalog.insert_object(3, vec![footer_entry(region_id, 20, 30)]), + "are not strictly increasing", + ); + assert_corrupted( + catalog.insert_object(3, vec![footer_entry(region_id, 5, 9)]), + "are not strictly increasing", + ); + assert_corrupted( + catalog.insert_object(1, vec![footer_entry(region_id, 15, 19)]), + "are not strictly increasing", + ); + assert_eq!(1, catalog.objects_in_order().count()); + } + + #[test] + fn test_catalog_rejects_duplicate_region_and_invalid_ranges() { + let region_id = RegionId::new(1, 1); + let mut catalog = ObjectCatalog::default(); + assert_corrupted( + catalog.insert_object( + 1, + vec![footer_entry(region_id, 1, 1), footer_entry(region_id, 2, 2)], + ), + "duplicate footer entries", + ); + assert_corrupted( + catalog.insert_object(1, vec![]), + "object 1 has an empty footer", + ); + + let mut invalid = footer_entry(region_id, 2, 1); + invalid.entry_count = 0; + assert_corrupted( + catalog.insert_object(1, vec![invalid]), + "has invalid entry range 2..=1", + ); + + let error = catalog + .objects_for_entry_range(region_id, 2, 1) + .unwrap_err(); + assert!( + matches!(error, Error::InvalidWalEntryRange { start_entry_id, end_entry_id, .. } if start_entry_id == 2 && end_entry_id == 1), + "unexpected error: {error:?}" + ); + } + + fn footer_entry(region_id: RegionId, min_entry_id: u64, max_entry_id: u64) -> FooterEntry { + FooterEntry { + region_id, + min_entry_id, + max_entry_id, + entry_count: max_entry_id + .checked_sub(min_entry_id) + .and_then(|count| count.checked_add(1)) + .unwrap_or(0) as u32, + segment_offset: min_entry_id.wrapping_mul(100), + segment_len: 100, + segment_crc32: min_entry_id as u32, + } + } + + fn object_seqs(objects: &[(u64, &FooterEntry)]) -> Vec { + objects.iter().map(|(object_seq, _)| *object_seq).collect() + } + + fn assert_corrupted(result: Result<()>, expected_reason: &str) { + match result { + Err(Error::CorruptedWalObject { reason, .. }) => assert!( + reason.contains(expected_reason), + "expected reason to contain {expected_reason:?}, actual {reason:?}" + ), + other => panic!("expected a corrupted object error, actual {other:?}"), + } + } +} diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs new file mode 100644 index 00000000000..0af6ac70278 --- /dev/null +++ b/src/log-store/src/object_store_wal/io.rs @@ -0,0 +1,558 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Object store access for WAL objects: a deterministic key layout and +//! conditional creates that make a retry of the same object a no-op. + +use bytes::Bytes; +use object_store::{ErrorKind, ObjectStore}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::error::{ + CorruptedWalObjectSnafu, InvalidWalObjectStoreSnafu, Result, WalObjectConflictSnafu, + WalObjectStoreSnafu, +}; + +/// Width of the zero-padded object sequence in an object key, wide enough for +/// [`u64::MAX`] so that lexicographic and numeric order agree. +const OBJECT_SEQ_WIDTH: usize = 20; +const OBJECT_SUFFIX: &str = ".wal"; + +/// Outcome of a conditional create. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PutResult { + /// This call created the object. + Created, + /// The object already held the same content, so the call was a retry. + AlreadyPresent, +} + +/// An object discovered by a LIST. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ListedObject { + pub(crate) object_seq: u64, + pub(crate) path: String, + /// Length of the object in bytes, as reported by the listing. + pub(crate) size: u64, +} + +/// Reads and writes the WAL objects under one prefix. +pub(crate) struct ObjectStoreIo { + store: ObjectStore, + object_prefix: String, +} + +impl ObjectStoreIo { + /// Binds `store` to the objects under `prefix`. The store must support + /// conditional creates, which is what keeps a retried write from + /// overwriting a durable object. + pub(crate) fn new(store: ObjectStore, prefix: impl AsRef) -> Result { + let object_prefix = normalize_prefix(prefix.as_ref())?; + ensure!( + store.info().capability().write_with_if_not_exists, + InvalidWalObjectStoreSnafu { + reason: "object store does not support conditional create", + } + ); + Ok(Self { + store, + object_prefix, + }) + } + + /// Creates the object `object_seq` unless it already exists. A retry with + /// the same content is accepted, different content is rejected. + /// + /// When the store reports that the object exists but the read that + /// compares its content fails, the read failure is returned, so that its + /// retry hint says whether comparing again may settle the outcome. Any + /// other create failure is returned as the write failure it is. + pub(crate) async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + let path = self.object_path(object_seq); + let write_result = self + .store + .write_with(&path, content.clone()) + .if_not_exists(true) + .await; + + match write_result { + Ok(_) => Ok(PutResult::Created), + Err(error) => match self.store.read(&path).await { + Ok(existing) if existing.to_bytes() == content => Ok(PutResult::AlreadyPresent), + Ok(_) => WalObjectConflictSnafu { path }.fail(), + Err(read_error) if reports_existing_object(&error) => { + Err(read_error).context(WalObjectStoreSnafu { + operation: "read", + path, + }) + } + Err(_) => Err(error).context(WalObjectStoreSnafu { + operation: "write", + path, + }), + }, + } + } + + /// Reads the object `object_seq`. + pub(crate) async fn get(&self, object_seq: u64) -> Result { + let path = self.object_path(object_seq); + self.store + .read(&path) + .await + .map(|content| content.to_bytes()) + .context(WalObjectStoreSnafu { + operation: "read", + path, + }) + } + + /// Reads `len` bytes of the object `object_seq` starting at `offset`. The + /// range must lie inside the object; one that reaches past its end fails. + pub(crate) async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + let path = self.object_path(object_seq); + let end = offset + .checked_add(len) + .with_context(|| CorruptedWalObjectSnafu { + reason: format!("byte range {offset}..{len} overflows the object"), + })?; + self.store + .read_with(&path) + .range(offset..end) + .await + .map(|content| content.to_bytes()) + .context(WalObjectStoreSnafu { + operation: "read", + path, + }) + } + + /// Lists the objects under the prefix, ordered by object sequence. Keys + /// that do not follow the object layout are ignored. + pub(crate) async fn list(&self) -> Result> { + let entries = self + .store + .list(&self.object_prefix) + .await + .with_context(|_| WalObjectStoreSnafu { + operation: "list", + path: self.object_prefix.clone(), + })?; + let mut objects = entries + .into_iter() + .filter_map(|entry| { + self.parse_object_seq(entry.path()) + .map(|object_seq| ListedObject { + object_seq, + path: entry.path().to_string(), + size: entry.metadata().content_length(), + }) + }) + .collect::>(); + objects.sort_unstable_by_key(|object| object.object_seq); + Ok(objects) + } + + pub(crate) fn object_path(&self, object_seq: u64) -> String { + format!( + "{}{object_seq:0OBJECT_SEQ_WIDTH$}{OBJECT_SUFFIX}", + self.object_prefix + ) + } + + fn parse_object_seq(&self, path: &str) -> Option { + let value = path + .strip_prefix(&self.object_prefix)? + .strip_suffix(OBJECT_SUFFIX)?; + if value.len() != OBJECT_SEQ_WIDTH || !value.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + value.parse().ok() + } +} + +/// Returns true when a conditional create failed because the object exists. +/// Only the precondition failure says so: a store may report `AlreadyExists` +/// for an unrelated path, such as a parent that is a file. +fn reports_existing_object(error: &object_store::Error) -> bool { + error.kind() == ErrorKind::ConditionNotMatch +} + +fn normalize_prefix(prefix: &str) -> Result { + let prefix = prefix.trim(); + ensure!( + !prefix.is_empty() && !prefix.starts_with('/'), + InvalidWalObjectStoreSnafu { + reason: format!("object prefix {prefix:?} is empty or absolute"), + } + ); + let prefix = prefix.strip_suffix('/').unwrap_or(prefix); + ensure!( + !prefix + .split('/') + .any(|segment| segment.is_empty() || segment == "." || segment == ".."), + InvalidWalObjectStoreSnafu { + reason: format!("object prefix {prefix:?} has an empty or relative component"), + } + ); + Ok(format!("{prefix}/objects/")) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_error::ext::{ErrorExt, RetryHint}; + use common_test_util::temp_dir::create_temp_dir; + use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use object_store::secure_fs::SecureFsRoot; + use object_store::services::Memory; + + use super::*; + use crate::error::Error; + + fn memory_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap() + } + + /// A reader whose every read fails with a temporary error. + struct FailingReader; + + impl oio::Read for FailingReader { + async fn open( + &self, + _range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, Box)> { + Err(injected_failure("read").set_temporary()) + } + + async fn read( + &self, + _range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, mock::Buffer)> { + Err(injected_failure("read").set_temporary()) + } + } + + /// A writer that fails to close, so nothing is written. + struct FailingWriter; + + impl oio::Write for FailingWriter { + async fn write(&mut self, _buffer: mock::Buffer) -> mock::Result<()> { + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + Err(injected_failure("write")) + } + + async fn abort(&mut self) -> mock::Result<()> { + Ok(()) + } + } + + fn injected_failure(operation: &str) -> mock::Error { + mock::Error::new( + mock::ErrorKind::Unexpected, + format!("injected {operation} failure"), + ) + } + + fn with_failing_reads(store: ObjectStore) -> ObjectStore { + store.layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(|_, _, _| Box::new(FailingReader))) + .build() + .unwrap(), + ) + } + + fn with_failing_writes(store: ObjectStore) -> ObjectStore { + store.layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(|_, _, _| Box::new(FailingWriter))) + .build() + .unwrap(), + ) + } + + fn memory_io() -> ObjectStoreIo { + ObjectStoreIo::new(memory_store(), "datanodes/1/epochs/2/").unwrap() + } + + #[test] + fn test_io_builds_and_parses_deterministic_object_paths() { + let io = ObjectStoreIo::new(memory_store(), " datanodes/1/epochs/2/ ").unwrap(); + assert_eq!( + "datanodes/1/epochs/2/objects/00000000000000000000.wal", + io.object_path(0) + ); + assert_eq!( + "datanodes/1/epochs/2/objects/00000000000000000042.wal", + io.object_path(42) + ); + assert_eq!( + "datanodes/1/epochs/2/objects/18446744073709551615.wal", + io.object_path(u64::MAX) + ); + assert_eq!(Some(42), io.parse_object_seq(&io.object_path(42))); + assert_eq!( + None, + io.parse_object_seq("datanodes/1/epochs/2/objects/42.wal") + ); + } + + #[test] + fn test_io_rejects_invalid_prefixes() { + for prefix in [ + "", + " ", + "/absolute", + ".", + "./node", + "node/./epoch", + "..", + "node/../epoch", + "node//epoch", + "node//", + "node///", + ] { + match ObjectStoreIo::new(memory_store(), prefix) { + Err(Error::InvalidWalObjectStore { .. }) => {} + Err(error) => panic!("unexpected error for prefix {prefix:?}: {error:?}"), + Ok(_) => panic!("expected prefix {prefix:?} to be rejected"), + } + } + } + + #[tokio::test] + async fn test_io_puts_gets_and_lists_objects_by_sequence() { + let io = memory_io(); + + for object_seq in [10, 2, 1] { + assert_eq!( + PutResult::Created, + io.put_if_absent(object_seq, Bytes::from(object_seq.to_string())) + .await + .unwrap() + ); + } + + assert_eq!(Bytes::from_static(b"2"), io.get(2).await.unwrap()); + let objects = io.list().await.unwrap(); + assert_eq!(vec![1, 2, 10], object_seqs(objects.clone())); + assert_eq!(io.object_path(1), objects[0].path); + assert_eq!( + vec![1, 1, 2], + objects.iter().map(|object| object.size).collect::>() + ); + } + + #[tokio::test] + async fn test_io_get_range_reads_a_slice_inside_the_object() { + let io = memory_io(); + io.put_if_absent(3, Bytes::from_static(b"0123456789")) + .await + .unwrap(); + + assert_eq!( + Bytes::from_static(b"234"), + io.get_range(3, 2, 3).await.unwrap() + ); + assert_eq!( + Bytes::from_static(b"89"), + io.get_range(3, 8, 2).await.unwrap() + ); + for (object_seq, offset, len) in [(3, 8, 5), (3, 10, 1), (7, 0, 1)] { + let error = io.get_range(object_seq, offset, len).await.unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(object_seq)), + "unexpected error for range {offset}..{len} of object {object_seq}: {error:?}" + ); + } + let error = io.get_range(3, u64::MAX, 1).await.unwrap_err(); + assert!( + matches!(error, Error::CorruptedWalObject { .. }), + "unexpected error: {error:?}" + ); + } + + #[tokio::test] + async fn test_io_same_content_is_an_idempotent_retry() { + let io = memory_io(); + let content = Bytes::from_static(b"immutable"); + + assert_eq!( + PutResult::Created, + io.put_if_absent(7, content.clone()).await.unwrap() + ); + assert_eq!( + PutResult::AlreadyPresent, + io.put_if_absent(7, content.clone()).await.unwrap() + ); + assert_eq!(content, io.get(7).await.unwrap()); + } + + #[tokio::test] + async fn test_io_rejects_different_content_without_overwriting() { + let io = memory_io(); + let original = Bytes::from_static(b"original"); + io.put_if_absent(7, original.clone()).await.unwrap(); + + let error = io + .put_if_absent(7, Bytes::from_static(b"replacement")) + .await + .unwrap_err(); + + assert!( + matches!(error, Error::WalObjectConflict { ref path, .. } if path == &io.object_path(7)), + "unexpected error: {error:?}" + ); + assert_eq!(original, io.get(7).await.unwrap()); + } + + #[tokio::test] + async fn test_io_retry_of_an_existing_object_reports_the_failed_read() { + let store = memory_store(); + let prefix = "datanodes/1/epochs/2/"; + let content = Bytes::from_static(b"immutable"); + ObjectStoreIo::new(store.clone(), prefix) + .unwrap() + .put_if_absent(7, content.clone()) + .await + .unwrap(); + + // The object exists, so the create is refused, and the read that would + // settle the retry fails: that failure and its retry hint come back. + let io = ObjectStoreIo::new(with_failing_reads(store.clone()), prefix).unwrap(); + let error = io.put_if_absent(7, content.clone()).await.unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(7)), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + + // A create that fails for another reason is reported as the write + // failure it is, although the read of the missing object fails too. + let io = ObjectStoreIo::new(with_failing_writes(store), prefix).unwrap(); + let error = io.put_if_absent(8, content).await.unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { operation: "write", ref path, .. } if path == &io.object_path(8)), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::NonRetryable, error.retry_hint()); + } + + #[tokio::test] + async fn test_io_create_under_a_file_parent_is_reported_as_the_write_failure() { + let temp_dir = create_temp_dir("object_store_wal_io_file_parent"); + std::fs::write(temp_dir.path().join("datanodes"), []).unwrap(); + let store = SecureFsRoot::open(temp_dir.path()) + .unwrap() + .build_operator(); + let io = ObjectStoreIo::new(store, "datanodes/1/epochs/2/").unwrap(); + + // The store reports `AlreadyExists` for the parent, not for the + // object, so the create failure itself comes back. + let error = io + .put_if_absent(7, Bytes::from_static(b"wal")) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { operation: "write", ref path, .. } if path == &io.object_path(7)), + "unexpected error: {error:?}" + ); + } + + #[tokio::test] + async fn test_io_get_of_a_missing_object_fails() { + let io = memory_io(); + + let error = io.get(7).await.unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(7)), + "unexpected error: {error:?}" + ); + } + + #[tokio::test] + async fn test_io_list_ignores_objects_outside_the_wal_layout() { + let io = memory_io(); + io.store.write("unrelated", "outside").await.unwrap(); + io.store + .write( + "datanodes/1/epochs/2/objects/not-a-sequence.wal", + "malformed", + ) + .await + .unwrap(); + io.put_if_absent(3, Bytes::from_static(b"wal")) + .await + .unwrap(); + + assert_eq!(vec![3], object_seqs(io.list().await.unwrap())); + } + + #[tokio::test] + async fn test_io_list_is_isolated_by_prefix() { + let store = memory_store(); + let first = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/1").unwrap(); + let second = ObjectStoreIo::new(store, "datanodes/1/epochs/2").unwrap(); + first + .put_if_absent(1, Bytes::from_static(b"first")) + .await + .unwrap(); + second + .put_if_absent(2, Bytes::from_static(b"second")) + .await + .unwrap(); + + assert_eq!(vec![1], object_seqs(first.list().await.unwrap())); + assert_eq!(vec![2], object_seqs(second.list().await.unwrap())); + } + + #[tokio::test] + async fn test_io_equal_sequences_in_different_prefixes_do_not_conflict() { + let store = memory_store(); + let first = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/1").unwrap(); + let second = ObjectStoreIo::new(store, "datanodes/1/epochs/2").unwrap(); + + assert_eq!( + PutResult::Created, + first + .put_if_absent(7, Bytes::from_static(b"first")) + .await + .unwrap() + ); + assert_eq!( + PutResult::Created, + second + .put_if_absent(7, Bytes::from_static(b"different")) + .await + .unwrap() + ); + assert_eq!(Bytes::from_static(b"first"), first.get(7).await.unwrap()); + assert_eq!( + Bytes::from_static(b"different"), + second.get(7).await.unwrap() + ); + } + + fn object_seqs(objects: Vec) -> Vec { + objects + .into_iter() + .map(|object| object.object_seq) + .collect() + } +} From b1106a9c5a7ab4d3a35932e85d9761af830f4c7c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 18 Sep 2026 07:40:33 +0000 Subject: [PATCH 05/63] perf(promql): propagate matching-label filters between binary operands (#9202) * perf(promql): propagate matching-label filters between binary operands A one-to-one arithmetic binary expression inner-joins its operands on the matching labels, so every row that survives the join already satisfies the other operand's equality matchers on those labels. Copy those matchers to the other operand so both scans drop non-joining series before execution instead of feeding them to the join. Both operands are planned before the rewrite: a selector matcher can also constrain a value field, and only the planned contexts tell tags and fields apart. A matcher is copied only when its name is a tag column on both sides, its value is non-empty, and it is one of the matching labels. Operands are limited to vector selectors, parentheses, and label-preserving range functions applied directly to a matrix selector; `ignoring(...)`, group modifiers, fill values, set and comparison operators, regex matchers and `or` matcher groups are left alone. Selector matchers are now deduplicated in place instead of through a `HashSet`, so the generated scan filter keeps a stable order once a selector carries more than one matcher. Signed-off-by: Dennis Zhuang * perf(promql): propagate ignoring, non-equality and aggregated matchers Widens the matcher propagation added in the previous commit to the shapes it was leaving on the table. The join compares matching labels with plain column equality (`normalized_match_key_expr` and its coalescing are confined to `or`), so any predicate on a matching label is already enforced on both sides for every surviving pair: it originates on one operand, and the join carries it to the pairs it forms. Copying it to the other operand can only drop rows that had no surviving partner. That argument does not depend on the matcher kind, so regular expressions, negations and empty values now propagate too. `=~".*"` stays out: it lowers to no filter at all, so copying it would only force a re-plan. `ignoring(...)` is no longer rejected. A label is a join key exactly when it is a tag on both sides and not named in `ignoring`, which is what `binary_join_key_columns` computes and what the caller can now answer from the two planned contexts. Operands may now be aggregations that partition by their grouping labels, which is the shape most real queries use. `agg_modifier_to_col` rewrites `ctx.tag_columns` to the grouping labels, so a label found in an aggregated operand's tags is a group key, and filtering the aggregate's input by it drops exactly the corresponding output groups. `topk`, `bottomk` and `limitk` are excluded because they select across a group and carry input labels through -- `prom_topk_bottomk_to_plan` leaves `ctx.tag_columns` alone, so the tag check cannot catch them. `count_values` is excluded because it adds an output label that does not exist in its input. The rollup whitelist now covers every label-preserving range function the planner implements, and finds the matrix argument by position so multi-argument rollups such as `quantile_over_time` and `predict_linear` qualify. `absent_over_time` stays out: it synthesizes a series from the matchers when its input has none. Re-running the sqlness case against an unmodified planner produces a byte-identical `.result` for all 27 queries. Signed-off-by: Dennis Zhuang * test(promql): pin that excluded matching labels stay on their own operand `ignoring(device)` and `on(host)` both leave `device` out of the join keys, so a `device` matcher must not reach the other operand. Neither case was covered end to end: a regression there silently drops the left operand's `eth1` series instead of returning them, which the two added queries now catch. Also corrects the comment at the rewrite site, which still described re-planning as touching a leaf selector. Aggregated operands are re-planned as a whole; what holds is that the rewrite only ever adds matchers to a selector, so the enclosing operand's table reference, time index and field columns are unchanged. Signed-off-by: Dennis Zhuang * refactor(promql): log why matching-filter propagation was skipped Splits the guard chain into `try_propagate`, whose `Err` names the reason the rewrite does not apply, and logs it once in `propagate`. Debugging a query that did not get the filter no longer means stepping through the guards. The reasons also replace the comments that used to explain the same conditions, and the remaining comments lose the parts that restated the code or repeated each other. Two of them were wrong rather than verbose. Saying `topk`'s input "must not be filtered" reads as a claim about PromQL: `topk(1, m{host="x"})` is perfectly legal, and what matters is that filtering before `topk` changes the candidate set it ranks. Saying the subset-matching case "keeps its many-to-many result" read as an endorsement of behaviour Prometheus rejects outright; it now records that the query is a cross product here and points at #9209. The propagation cannot change whether such a check would fire: series in one match group agree on every matching label, so a matcher over one of those labels keeps all of them or drops all of them. Signed-off-by: Dennis Zhuang * fix(promql): keep copied matchers out of the operand's selector metadata Re-planning an operand with a copied matcher also rebuilt its `PromPlannerContext::selector_matcher`, and that context is what an enclosing expression reads. `create_absent_plan` turns the equality matchers found there into the labels `absent()` reports, so absent(counter_metric{host="missing"} / on(host, device) gauge_metric) gained a `host="missing"` label that `main` does not produce. The inner expression was empty either way; only the reported label set changed. The copied matcher belongs to the scan, not to the operand's identity, so both re-planned contexts now keep the matchers their operand was written with. `selector_matcher` has one other reader, `create_table_scan_plan`, which consumes it while the operand is being planned and is unaffected. Reported by @discord9. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/planner.rs | 124 +++++- .../src/promql/planner/matching_filters.rs | 356 ++++++++++++++++ .../common/tql/matching_filter.result | 389 ++++++++++++++++++ .../standalone/common/tql/matching_filter.sql | 137 ++++++ 4 files changed, 997 insertions(+), 9 deletions(-) create mode 100644 src/query/src/promql/planner/matching_filters.rs create mode 100644 tests/cases/standalone/common/tql/matching_filter.result create mode 100644 tests/cases/standalone/common/tql/matching_filter.sql diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c51cdb05831..17cc36a41da 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod matching_filters; + use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::sync::Arc; use std::time::UNIX_EPOCH; @@ -1490,21 +1492,53 @@ impl PromPlanner { } // both are columns. join them on time index (None, None) => { - let left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?; + let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?; let left_field_columns = self.ctx.field_columns.clone(); let left_time_index_column = self.ctx.time_index_column.clone(); let mut left_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let left_context = self.ctx.clone(); + let mut left_context = self.ctx.clone(); - let right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?; + let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?; let right_field_columns = self.ctx.field_columns.clone(); let right_time_index_column = self.ctx.time_index_column.clone(); let mut right_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let right_context = self.ctx.clone(); + let mut right_context = self.ctx.clone(); + + // Both operands are planned first because only the planned contexts tell a tag + // from a value field. The rewrite merely adds matchers to a selector, so the + // table reference, time index and field columns captured above stay valid. + if let Some(rewritten) = matching_filters::propagate( + binary_expr, + &left_context.tag_columns, + &right_context.tag_columns, + ) { + // A copied matcher belongs to the scan, not to the operand's identity: + // `absent()` turns `selector_matcher` into the labels it reports, so the + // re-planned context keeps the matchers the operand was written with. + if rewritten.lhs.as_ref() != lhs.as_ref() { + let selectors = std::mem::take(&mut left_context.selector_matcher); + left_input = self + .prom_expr_to_plan(&rewritten.lhs, query_engine_state) + .await?; + left_context = self.ctx.clone(); + left_context.selector_matcher = selectors; + } + if rewritten.rhs.as_ref() != rhs.as_ref() { + let selectors = std::mem::take(&mut right_context.selector_matcher); + right_input = self + .prom_expr_to_plan(&rewritten.rhs, query_engine_state) + .await?; + right_context = self.ctx.clone(); + right_context.selector_matcher = selectors; + } + // The code below reads `self.ctx` as the right operand's context. + self.ctx = right_context.clone(); + } + let left_is_empty_metric = Self::is_empty_metric(&left_input); let right_is_empty_metric = Self::is_empty_metric(&right_input); @@ -2344,7 +2378,9 @@ impl PromPlanner { self.ctx.table_name = metric_name; - let mut matchers = HashSet::new(); + // Deduplicate in place instead of through a `HashSet`: the scan filter is built in + // this order, and a hashed order makes the plan of one query vary between runs. + let mut matchers: Vec = Vec::with_capacity(label_matchers.matchers.len()); for matcher in &label_matchers.matchers { // TODO(ruihang): support other metric match ops if matcher.name == FIELD_COLUMN_MATCHER { @@ -2363,11 +2399,13 @@ impl PromPlanner { self.ctx.schema_name = Some(matcher.value.clone()); } else if matcher.name != METRIC_NAME { self.ctx.selector_matcher.push(matcher.clone()); - let _ = matchers.insert(matcher.clone()); + if !matchers.contains(matcher) { + matchers.push(matcher.clone()); + } } } - Ok(Matchers::new(matchers.into_iter().collect())) + Ok(Matchers::new(matchers)) } async fn selector_to_series_normalize_plan( @@ -10686,13 +10724,13 @@ mod test { \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", ); @@ -13309,6 +13347,74 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom assert!(!contains_histogram_fold(&plan), "{plan:?}"); } + async fn build_matching_filter_plan(query: &str) -> String { + let table_provider = build_test_table_provider_with_distinct_tags(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() + } + + #[tokio::test] + async fn binary_matching_label_filter_reaches_both_operands() { + for query in [ + r#"metric_a / metric_b{host="foo"}"#, + r#"metric_a / on(host, device) metric_b{host="foo"}"#, + r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, + r#"metric_a / ignoring(device) metric_b{host="foo"}"#, + r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, + r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } + } + + #[tokio::test] + async fn binary_matching_label_filter_skips_selecting_aggregations() { + // `topk` ranks its input, so filtering before it changes the candidate set. + let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); + } + + #[tokio::test] + async fn binary_value_field_matcher_stays_on_its_own_operand() { + let value = greptime_value(); + for query in [ + format!(r#"metric_a / metric_b{{{value}="2"}}"#), + format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), + ] { + let plan = build_matching_filter_plan(&query).await; + assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); + } + } + + #[tokio::test] + async fn binary_matching_label_filter_skips_unproven_expressions() { + for query in [ + r#"metric_a > on(host, device) metric_b{host="foo"}"#, + r#"metric_a / on(host) group_left metric_b{host="foo"}"#, + r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, + r#"metric_a / on(device) metric_b{host="foo"}"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); + } + } + #[tokio::test] async fn test_or_context_uses_left_qualified_output() { let case = r#"(normal_metric or other_metric) + 1"#; diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs new file mode 100644 index 00000000000..3ed5ee93baf --- /dev/null +++ b/src/query/src/promql/planner/matching_filters.rs @@ -0,0 +1,356 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Propagation of label matchers between the operands of a PromQL binary expression. + +use promql_parser::label::{MatchOp, Matcher, Matchers}; +use promql_parser::parser::{ + AggregateExpr, BinaryExpr, Expr, LabelModifier, VectorMatchCardinality, token, +}; + +/// [`super::PromPlanner::matchers_to_expr`] emits no filter for this, so copying it would only +/// force a re-plan. +const MATCH_ALL_REGEX: &str = "^(?:.*)$"; + +/// Rollup functions emitting at most one output series per input series, carrying the labels of +/// the matrix selector they wrap. `absent_over_time` is excluded: it synthesizes a series from +/// the matchers when its input has none. +const LABEL_PRESERVING_RANGE_FUNCTIONS: [&str; 20] = [ + "avg_over_time", + "changes", + "count_over_time", + "delta", + "deriv", + "double_exponential_smoothing", + "idelta", + "increase", + "irate", + "last_over_time", + "max_over_time", + "min_over_time", + "predict_linear", + "present_over_time", + "quantile_over_time", + "rate", + "resets", + "stddev_over_time", + "stdvar_over_time", + "sum_over_time", +]; + +/// Copies label matchers from one operand of `binary` to the other, so both scans discard +/// non-joining series instead of the join. +/// +/// One-to-one arithmetic operands are inner-joined on their matching labels with plain column +/// equality, so a matcher on a matching label already holds for every surviving pair. Copying it +/// to the other operand can only drop rows that had no partner, whatever the matcher kind, and +/// `NULL` pairs only with `NULL`. Nor can it split a match group, whose series agree on every +/// matching label, so it does not affect a cardinality check such as #9209. +/// +/// `left_tags` and `right_tags` are the tag columns of the planned operands: a matcher may just +/// as well constrain a value field, which the parsed expression does not distinguish from a +/// label. For an aggregated operand they are its grouping labels, which is what makes filtering +/// its input equivalent to filtering its output. +pub(super) fn propagate( + binary: &BinaryExpr, + left_tags: &[String], + right_tags: &[String], +) -> Option { + match try_propagate(binary, left_tags, right_tags) { + Ok(rewritten) => Some(rewritten), + Err(reason) => { + common_telemetry::debug!("Matching filter not propagated ({reason}): {binary}"); + None + } + } +} + +/// `Err` carries why the rewrite does not apply, for [`propagate`] to log. +fn try_propagate( + binary: &BinaryExpr, + left_tags: &[String], + right_tags: &[String], +) -> Result { + if !matches!( + binary.op.id(), + token::T_ADD | token::T_SUB | token::T_MUL | token::T_DIV | token::T_MOD | token::T_POW + ) { + return Err("operator is not arithmetic"); + } + if let Some(modifier) = &binary.modifier { + if !matches!(modifier.card, VectorMatchCardinality::OneToOne) { + return Err("matching is not one-to-one"); + } + if modifier.fill_values.lhs.is_some() || modifier.fill_values.rhs.is_some() { + return Err("operand carries a fill modifier"); + } + } + let matching = binary + .modifier + .as_ref() + .and_then(|modifier| modifier.matching.as_ref()); + + let mut rewritten = binary.clone(); + let left = selector_matchers(&mut rewritten.lhs).ok_or("left operand is not a selector")?; + let right = selector_matchers(&mut rewritten.rhs).ok_or("right operand is not a selector")?; + if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { + return Err("selector has an or matcher group"); + } + + let is_matching_label = |name: &String| { + // PromQL reserves the `__` prefix; none of those names is a join key. + !name.starts_with("__") + && left_tags.contains(name) + && right_tags.contains(name) + && match matching { + None => true, + Some(LabelModifier::Include(on)) => on.labels.contains(name), + Some(LabelModifier::Exclude(ignoring)) => !ignoring.labels.contains(name), + } + }; + let constraints = left + .matchers + .iter() + .chain(&right.matchers) + .filter(|matcher| is_matching_label(&matcher.name) && !matches_every_value(matcher)) + .cloned() + .collect::>(); + + let mut changed = false; + for matcher in constraints { + for target in [&mut left.matchers, &mut right.matchers] { + if !target.contains(&matcher) { + target.push(matcher.clone()); + changed = true; + } + } + } + if !changed { + return Err("operands already carry the same matching-label matchers"); + } + Ok(rewritten) +} + +fn matches_every_value(matcher: &Matcher) -> bool { + matches!(&matcher.op, MatchOp::Re(re) if re.as_str() == MATCH_ALL_REGEX) +} + +/// Aggregations partitioning their input by the grouping labels, so that dropping input series +/// by a grouping label drops exactly the matching output series and leaves the remaining +/// aggregated values untouched. +/// +/// `topk`, `bottomk` and `limitk` rank across a group and carry the input labels through, so +/// filtering before them changes the candidate set. `count_values` adds an output label that +/// does not exist in its input. +fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { + matches!( + aggregate.op.id(), + token::T_SUM + | token::T_AVG + | token::T_COUNT + | token::T_MIN + | token::T_MAX + | token::T_GROUP + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE + ) +} + +/// Returns the matchers of the vector selector an operand scans, or `None` when the operand's +/// output labels are not proven to be that selector's labels. +fn selector_matchers(expr: &mut Expr) -> Option<&mut Matchers> { + match expr { + Expr::VectorSelector(selector) => Some(&mut selector.matchers), + Expr::Paren(paren) => selector_matchers(&mut paren.expr), + Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { + selector_matchers(&mut aggregate.expr) + } + Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => { + // Every other argument is a scalar, so position does not matter. + let matrix = single_matrix_argument(&call.args.args)?; + match call.args.args[matrix].as_mut() { + Expr::MatrixSelector(selector) => Some(&mut selector.vs.matchers), + _ => None, + } + } + // Label rewriting and subqueries need separate proofs. + _ => None, + } +} + +fn single_matrix_argument(args: &[Box]) -> Option { + let mut found = None; + for (index, arg) in args.iter().enumerate() { + if matches!(arg.as_ref(), Expr::MatrixSelector(_)) { + if found.is_some() { + return None; + } + found = Some(index); + } + } + found +} + +#[cfg(test)] +mod tests { + use promql_parser::parser::parse; + + use super::*; + + fn rewrite_with(query: &str, left_tags: &[&str], right_tags: &[&str]) -> Option { + let Expr::Binary(binary) = parse(query).unwrap() else { + panic!("expected binary") + }; + let owned = |tags: &[&str]| tags.iter().map(|tag| tag.to_string()).collect::>(); + propagate(&binary, &owned(left_tags), &owned(right_tags)).map(Expr::Binary) + } + + fn rewrite(query: &str) -> Option { + rewrite_with(query, &["host", "zone"], &["host", "zone"]) + } + + #[track_caller] + fn assert_rewrite(query: &str, expected: &str) { + assert_eq!(rewrite(query).unwrap(), parse(expected).unwrap(), "{query}"); + // The rewrite adds every constraint to both sides, so re-running it is a no-op. + assert!(rewrite(expected).is_none(), "{expected}"); + } + + #[test] + fn propagates_only_explicit_matching_labels() { + assert_rewrite( + r#"a / on(host) b{host="x",zone="y"}"#, + r#"a{host="x"} / on(host) b{host="x",zone="y"}"#, + ); + } + + #[test] + fn propagates_labels_not_named_in_ignoring() { + assert_rewrite( + r#"a{zone="y"} / ignoring(zone) b{host="x"}"#, + r#"a{zone="y",host="x"} / ignoring(zone) b{host="x"}"#, + ); + } + + #[test] + fn propagates_default_matching_without_metric_name() { + assert_rewrite( + r#"a / {host="x",__name__="b"}"#, + r#"a{host="x"} / {host="x",__name__="b"}"#, + ); + assert!(rewrite(r#"a / {__name__="b"}"#).is_none()); + } + + #[test] + fn propagates_only_names_that_are_tags_on_both_sides() { + // `value` is a field column of both metrics, even when named in `on(...)`. + for query in [r#"a / b{value="2"}"#, r#"a / on(host,value) b{value="2"}"#] { + assert!(rewrite(query).is_none(), "{query}"); + } + assert_rewrite( + r#"a / b{host="x",value="2"}"#, + r#"a{host="x"} / b{host="x",value="2"}"#, + ); + // `zone` is a tag of the right metric only. + assert!(rewrite_with(r#"a / b{zone="y"}"#, &["host"], &["host", "zone"]).is_none()); + } + + #[test] + fn propagates_matchers_the_join_enforces_anyway() { + for (query, expected) in [ + ( + r#"a / b{host=~"x.*"}"#, + r#"a{host=~"x.*"} / b{host=~"x.*"}"#, + ), + (r#"a / b{host!="x"}"#, r#"a{host!="x"} / b{host!="x"}"#), + ( + r#"a / b{host!~"x.*"}"#, + r#"a{host!~"x.*"} / b{host!~"x.*"}"#, + ), + (r#"a / b{host=""}"#, r#"a{host=""} / b{host=""}"#), + ] { + assert_rewrite(query, expected); + } + assert!(rewrite(r#"a / b{host=~".*"}"#).is_none()); + } + + #[test] + fn propagates_through_partitioning_aggregations() { + assert_rewrite( + r#"sum by(host) (a) / on(host) max by(host) (b{host="x"})"#, + r#"sum by(host) (a{host="x"}) / on(host) max by(host) (b{host="x"})"#, + ); + assert_rewrite( + r#"avg without(zone) (rate(a[5m])) / b{host="x"}"#, + r#"avg without(zone) (rate(a{host="x"}[5m])) / b{host="x"}"#, + ); + } + + #[test] + fn leaves_selecting_aggregations_alone() { + for query in [ + r#"topk(3, a) / on(host) b{host="x"}"#, + r#"bottomk(3, a) / on(host) b{host="x"}"#, + r#"count_values("v", a) / on(host) b{host="x"}"#, + ] { + assert!(rewrite(query).is_none(), "{query}"); + } + } + + #[test] + fn finds_the_matrix_argument_of_multi_argument_rollups() { + assert_rewrite( + r#"quantile_over_time(0.9, a[5m]) / on(host) b{host="x"}"#, + r#"quantile_over_time(0.9, a{host="x"}[5m]) / on(host) b{host="x"}"#, + ); + assert_rewrite( + r#"predict_linear(a[5m], 60) / on(host) b{host="x"}"#, + r#"predict_linear(a{host="x"}[5m], 60) / on(host) b{host="x"}"#, + ); + assert_rewrite( + r#"double_exponential_smoothing(a[5m], 0.5, 0.5) / on(host) b{host="x"}"#, + r#"double_exponential_smoothing(a{host="x"}[5m], 0.5, 0.5) / on(host) b{host="x"}"#, + ); + } + + #[test] + fn preserves_windows_offsets_and_conflicting_matchers() { + assert_eq!( + rewrite( + r#"rate(a{host="x"}[5m] offset 1h) / on(host) count_over_time(b{host="y"}[1m])"# + ) + .unwrap(), + parse( + r#"rate(a{host="x",host="y"}[5m] offset 1h) / on(host) count_over_time(b{host="y",host="x"}[1m])"# + ) + .unwrap() + ); + } + + #[test] + fn leaves_unproven_semantics_unchanged() { + for query in [ + r#"a or on(host) b{host="x"}"#, + r#"a / on(host) group_left b{host="x"}"#, + r#"sum by(host)(a) / on(host) group_right b{host="x"}"#, + r#"label_replace(a,"host","x","zone",".*") / on(host) b{host="x"}"#, + r#"a > on(host) b{host="x"}"#, + r#"a / on(host) absent_over_time(b{host="x"}[5m])"#, + r#"a / on(host) (b{host="x"} + b)"#, + ] { + assert!(rewrite(query).is_none(), "{query}"); + } + } +} diff --git a/tests/cases/standalone/common/tql/matching_filter.result b/tests/cases/standalone/common/tql/matching_filter.result new file mode 100644 index 00000000000..98834ad0902 --- /dev/null +++ b/tests/cases/standalone/common/tql/matching_filter.result @@ -0,0 +1,389 @@ +create table gauge_metric( + ts timestamp time index, + host string, + device string, + val double, + primary key (host, device) +); + +Affected Rows: 0 + +insert into gauge_metric values + (0, 'host1', 'eth0', 2), + (0, 'host1', 'eth1', 2), + (0, 'host2', 'eth0', 2), + (5000, 'host1', 'eth0', 2), + (5000, 'host1', 'eth1', 2), + (5000, 'host2', 'eth0', 2); + +Affected Rows: 6 + +create table counter_metric( + ts timestamp time index, + host string, + device string, + val double, + primary key (host, device) +); + +Affected Rows: 0 + +insert into counter_metric values + (0, 'host1', 'eth0', 10), + (0, 'host1', 'eth1', 20), + (0, 'host2', 'eth0', 40), + (5000, 'host1', 'eth0', 10), + (5000, 'host1', 'eth1', 20), + (5000, 'host2', 'eth0', 40); + +Affected Rows: 6 + +-- Default matching. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / {__name__="gauge_metric", host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Explicit `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric{host="host1"} / on(host, device) gauge_metric{device="eth0"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Conflicting constraints on the same label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric{host="host1"} / on(host, device) gauge_metric{host="host2"}; + ++------+--------+----+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++------+--------+----+---------------------------------------+ ++------+--------+----+---------------------------------------+ + +-- `val` is a value field of both metrics, not a label, even when named in `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{val="2"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device, val) gauge_metric{val="2"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- A matcher on a name that is not a column of the metric. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{missing_label="x"}; + ++------+--------+----+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++------+--------+----+---------------------------------------+ ++------+--------+----+---------------------------------------+ + +-- An empty matcher value also matches an absent label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{device=""}; + ++------+--------+----+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++------+--------+----+---------------------------------------+ ++------+--------+----+---------------------------------------+ + +-- Range functions over a matrix selector, and offsets. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') count_over_time(counter_metric[10s]) / on(host, device) count_over_time(gauge_metric{host="host1"}[10s]); + ++-------+--------+---------------------+-----------------------------------------------------------------------------------------------------+ +| host | device | ts | counter_metric.prom_count_over_time(ts_range,val) / gauge_metric.prom_count_over_time(ts_range,val) | ++-------+--------+---------------------+-----------------------------------------------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 1.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 1.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 1.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 1.0 | ++-------+--------+---------------------+-----------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB +-- returns the cross product instead (#9209). Recorded here to show the rewrite does not +-- change it, not to endorse it. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- `ignoring(...)`: every label that is not ignored is a matching label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(missing_label) gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{host="host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- An ignored label is not a matching label: the left operand keeps its `eth1` series. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{device="eth0"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Same for a label left out of `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Matcher kinds other than equality are enforced by the join just the same. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host=~"host1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host!="host2"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{device!~"eth1"}; + ++-------+--------+---------------------+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++-------+--------+---------------------+---------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | +| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | ++-------+--------+---------------------+---------------------------------------+ + +-- Aggregations that partition by their grouping labels. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') sum by(host)(counter_metric) / on(host) sum by(host)(gauge_metric{host="host1"}); + ++-------+---------------------+-----------------------------------------------------------------------------+ +| host | ts | counter_metric.sum(counter_metric.val) / gauge_metric.sum(gauge_metric.val) | ++-------+---------------------+-----------------------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 7.5 | +| host1 | 1970-01-01T00:00:05 | 7.5 | ++-------+---------------------+-----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') sum by(host)(counter_metric{device="eth0"}) / on(host) sum by(host)(gauge_metric); + ++-------+---------------------+-----------------------------------------------------------------------------+ +| host | ts | counter_metric.sum(counter_metric.val) / gauge_metric.sum(gauge_metric.val) | ++-------+---------------------+-----------------------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 2.5 | +| host1 | 1970-01-01T00:00:05 | 2.5 | +| host2 | 1970-01-01T00:00:00 | 20.0 | +| host2 | 1970-01-01T00:00:05 | 20.0 | ++-------+---------------------+-----------------------------------------------------------------------------+ + +-- A label the aggregation drops is not a matching label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device="eth0"}); + ++-------+---------------------+---------------------------------------------------------+ +| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | ++-------+---------------------+---------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 10.0 | +| host1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | 1970-01-01T00:00:05 | 10.0 | +| host1 | 1970-01-01T00:00:05 | 5.0 | +| host2 | 1970-01-01T00:00:00 | 20.0 | +| host2 | 1970-01-01T00:00:05 | 20.0 | ++-------+---------------------+---------------------------------------------------------+ + +-- `topk` preserves its input labels, and filtering before `topk` changes the set of +-- candidates it ranks, so the rewrite must not reach its input. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') topk(1, counter_metric) / on(host, device) gauge_metric{host="host1"}; + ++------+--------+----+---------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.val | ++------+--------+----+---------------------------------------+ ++------+--------+----+---------------------------------------+ + +-- Operators and modifiers outside the rewritten subset. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric > on(host, device) gauge_metric{host="host1"}; + ++---------------------+------+-------+--------+ +| ts | val | host | device | ++---------------------+------+-------+--------+ +| 1970-01-01T00:00:00 | 10.0 | host1 | eth0 | +| 1970-01-01T00:00:00 | 20.0 | host1 | eth1 | +| 1970-01-01T00:00:05 | 10.0 | host1 | eth0 | +| 1970-01-01T00:00:05 | 20.0 | host1 | eth1 | ++---------------------+------+-------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) group_left sum by(host)(gauge_metric{host="host1"}); + ++-------+---------------------+---------------------------------------------------------+ +| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | ++-------+---------------------+---------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 2.5 | +| host1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | 1970-01-01T00:00:05 | 2.5 | +| host1 | 1970-01-01T00:00:05 | 5.0 | ++-------+---------------------+---------------------------------------------------------+ + +-- `absent()` reports the inner selector's equality matchers as labels, so a copied matcher +-- must not reach the context an enclosing expression reads. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') absent(counter_metric{host="missing"} / on(host, device) gauge_metric); + ++---------------------+---------------------------------------+ +| ts | counter_metric.val / gauge_metric.val | ++---------------------+---------------------------------------+ +| 1970-01-01T00:00:00 | 1.0 | +| 1970-01-01T00:00:05 | 1.0 | ++---------------------+---------------------------------------+ + +drop table gauge_metric; + +Affected Rows: 0 + +drop table counter_metric; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/tql/matching_filter.sql b/tests/cases/standalone/common/tql/matching_filter.sql new file mode 100644 index 00000000000..fa2bc3371d5 --- /dev/null +++ b/tests/cases/standalone/common/tql/matching_filter.sql @@ -0,0 +1,137 @@ +create table gauge_metric( + ts timestamp time index, + host string, + device string, + val double, + primary key (host, device) +); + +insert into gauge_metric values + (0, 'host1', 'eth0', 2), + (0, 'host1', 'eth1', 2), + (0, 'host2', 'eth0', 2), + (5000, 'host1', 'eth0', 2), + (5000, 'host1', 'eth1', 2), + (5000, 'host2', 'eth0', 2); + +create table counter_metric( + ts timestamp time index, + host string, + device string, + val double, + primary key (host, device) +); + +insert into counter_metric values + (0, 'host1', 'eth0', 10), + (0, 'host1', 'eth1', 20), + (0, 'host2', 'eth0', 40), + (5000, 'host1', 'eth0', 10), + (5000, 'host1', 'eth1', 20), + (5000, 'host2', 'eth0', 40); + +-- Default matching. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{host="host1"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / {__name__="gauge_metric", host="host1"}; + +-- Explicit `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host="host1"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric{host="host1"} / on(host, device) gauge_metric{device="eth0"}; + +-- Conflicting constraints on the same label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric{host="host1"} / on(host, device) gauge_metric{host="host2"}; + +-- `val` is a value field of both metrics, not a label, even when named in `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{val="2"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device, val) gauge_metric{val="2"}; + +-- A matcher on a name that is not a column of the metric. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / gauge_metric{missing_label="x"}; + +-- An empty matcher value also matches an absent label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{device=""}; + +-- Range functions over a matrix selector, and offsets. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') count_over_time(counter_metric[10s]) / on(host, device) count_over_time(gauge_metric{host="host1"}[10s]); + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{host="host1"}; + +-- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB +-- returns the cross product instead (#9209). Recorded here to show the rewrite does not +-- change it, not to endorse it. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; + +-- `ignoring(...)`: every label that is not ignored is a matching label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(missing_label) gauge_metric{host="host1"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{host="host1"}; + +-- An ignored label is not a matching label: the left operand keeps its `eth1` series. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{device="eth0"}; + +-- Same for a label left out of `on(...)`. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + +-- Matcher kinds other than equality are enforced by the join just the same. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host=~"host1"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{host!="host2"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host, device) gauge_metric{device!~"eth1"}; + +-- Aggregations that partition by their grouping labels. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') sum by(host)(counter_metric) / on(host) sum by(host)(gauge_metric{host="host1"}); + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') sum by(host)(counter_metric{device="eth0"}) / on(host) sum by(host)(gauge_metric); + +-- A label the aggregation drops is not a matching label. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device="eth0"}); + +-- `topk` preserves its input labels, and filtering before `topk` changes the set of +-- candidates it ranks, so the rewrite must not reach its input. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') topk(1, counter_metric) / on(host, device) gauge_metric{host="host1"}; + +-- Operators and modifiers outside the rewritten subset. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric > on(host, device) gauge_metric{host="host1"}; + +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') counter_metric / on(host) group_left sum by(host)(gauge_metric{host="host1"}); + +-- `absent()` reports the inner selector's equality matchers as labels, so a copied matcher +-- must not reach the context an enclosing expression reads. +-- SQLNESS SORT_RESULT 3 1 +tql eval(0, 5, '5s') absent(counter_metric{host="missing"} / on(host, device) gauge_metric); + +drop table gauge_metric; + +drop table counter_metric; From 09a9d9088da87b57bba0a6f0a1553598fc593636 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 18 Sep 2026 08:02:34 +0000 Subject: [PATCH 06/63] feat(otlp): preserve trace v2 events and links as JSON (follow-up to #9192) (#9232) feat(otlp): preserve trace v2 events and links as JSON Signed-off-by: luofucong --- src/servers/src/otlp/trace/v2.rs | 19 ++++--- tests-integration/src/otlp.rs | 57 ++++++++++++++++++- tests-integration/tests/http.rs | 14 ++++- .../cases/trace_v2_json2/case.toml | 2 +- .../cases/trace_v2_json2/setup.sql | 18 ++++-- .../cases/trace_v2_json2/verify.result | 23 +++++++- .../cases/trace_v2_json2/verify.sql | 9 +++ 7 files changed, 124 insertions(+), 18 deletions(-) diff --git a/src/servers/src/otlp/trace/v2.rs b/src/servers/src/otlp/trace/v2.rs index 7d2dfb40adb..6d4ee86e326 100644 --- a/src/servers/src/otlp/trace/v2.rs +++ b/src/servers/src/otlp/trace/v2.rs @@ -23,18 +23,17 @@ use crate::otlp::trace::v1::span_duration_nano; use crate::otlp::trace::{ DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, SCOPE_NAME_COLUMN, SCOPE_VERSION_COLUMN, SERVICE_NAME_COLUMN, - SPAN_ATTRIBUTES_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_COLUMN, SPAN_NAME_COLUMN, SPAN_STATUS_CODE, - SPAN_STATUS_MESSAGE_COLUMN, TIMESTAMP_COLUMN, TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, - TRACE_STATE_COLUMN, + SPAN_ATTRIBUTES_COLUMN, SPAN_EVENTS_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_COLUMN, SPAN_NAME_COLUMN, + SPAN_STATUS_CODE, SPAN_STATUS_MESSAGE_COLUMN, TIMESTAMP_COLUMN, TIMESTAMP_END_COLUMN, + TRACE_ID_COLUMN, TRACE_STATE_COLUMN, }; use crate::otlp::utils::{make_column_data, make_string_column_data}; use crate::row_writer::{self, MultiTableData, TableData}; // Preallocate for the fixed v2 schema: 3 timing columns, 3 span/trace IDs, // 2 kind/name columns, 2 status columns, 1 trace state, 2 scope name/version -// columns, 1 service name, and 3 JSON2 attribute columns (17 total). -// Events and links are excluded until ARRAY(JSON2) is supported. -const APPROXIMATE_COLUMN_COUNT: usize = 17; +// columns, 1 service name, 3 JSON2 attribute columns, and 2 JSON event/link columns. +const APPROXIMATE_COLUMN_COUNT: usize = 19; /// Converts trace spans into row insert requests for the main v2 trace table. pub(super) fn v2_to_grpc_main_insert_requests( @@ -139,7 +138,13 @@ fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { &mut row, )?; - // TODO(LFC): Store span_events and span_links once ARRAY(JSON2) is supported. + row_writer::write_json( + writer, + SPAN_EVENTS_COLUMN, + span.span_events.into(), + &mut row, + )?; + row_writer::write_json(writer, "span_links", span.span_links.into(), &mut row)?; writer.add_row(row); Ok(()) diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index fe4dfb8a9e5..dacab7ad6bd 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -84,6 +84,17 @@ mod test { "spanId": "9630f2916e2f7909", "name": "op", "kind": 2, "startTimeUnixNano": "1736480942444376000", "endTimeUnixNano": "1736480942444499000", + "events": [{ + "timeUnixNano": "1736480942444400000", + "name": "cache.hit", + "attributes": [{"key": "event.code", "value": {"intValue": "7"}}] + }], + "links": [{ + "traceId": "cc9e0991a2e63d274984bd44ee669203", + "spanId": "8f847259b0f6e1ab", + "traceState": "vendor=value", + "attributes": [{"key": "link.type", "value": {"stringValue": "follows_from"}}] + }], "attributes": [ {"key": "http.status_code", "value": {"intValue": "200"}}, {"key": "latency", "value": {"doubleValue": 12.5}}, @@ -147,6 +158,8 @@ mod test { "resource_attributes" JSON2( max_auto_expanded_paths = 100 ) NULL, + "span_events" JSON NULL, + "span_links" JSON NULL, TIME INDEX ("timestamp"), PRIMARY KEY ("service_name") ) @@ -199,7 +212,7 @@ WITH( .await?; assert_trace_v2_query( instance, - "SELECT COUNT(*) = 17 + "SELECT COUNT(*) = 19 FROM information_schema.columns WHERE table_name = 'trace_v2'", ctx.clone(), @@ -226,7 +239,22 @@ WITH( .await?; // Exercise each JSON2 column and expressions combining columns, before and after flush. + let event_time = + common_time::Timestamp::new_nanosecond(1736480942444400000).to_iso8601_string(); + let event_query = format!( + r#"SELECT COUNT(*) > 0 AND COUNT(*) = COUNT(CASE + WHEN json_get_string(span_events, '$[0].name') = 'cache.hit' + AND json_get_string(span_events, '$[0].time') = '{event_time}' + AND json_get_int(span_events, '$[0].attributes."event.code"') = 7 + AND json_get_string(span_links, '$[0].trace_id') = 'cc9e0991a2e63d274984bd44ee669203' + AND json_get_string(span_links, '$[0].span_id') = '8f847259b0f6e1ab' + AND json_get_string(span_links, '$[0].trace_state') = 'vendor=value' + AND json_get_string(span_links, '$[0].attributes."link.type"') = 'follows_from' + THEN 1 END) + FROM trace_v2"# + ); let json_queries = [ + event_query.as_str(), r#"SELECT COUNT(*) > 0 AND COUNT(*) = COUNT(CASE WHEN span_attributes.latency::DOUBLE * 2 = 25 AND upper(span_attributes.nested."a.b"[1]::STRING) = 'OK' @@ -358,7 +386,7 @@ WITH( ); assert_trace_v2_query( instance, - "SELECT COUNT(*) = 17 + "SELECT COUNT(*) = 19 FROM information_schema.columns WHERE table_name = 'trace_v2'", ctx.clone(), @@ -393,7 +421,7 @@ WITH( let error = instance .traces( instance.clone(), - request, + request.clone(), PipelineWay::OtlpTraceDirectV2, GreptimePipelineParams::default(), "trace_v1".to_string(), @@ -423,6 +451,29 @@ WITH( for sql in json_queries { assert_trace_v2_query(instance, sql, ctx.clone()).await?; } + // Empty events and links are stored as JSON arrays in the same 19-column table. + let mut empty_request = request; + let empty_span = &mut empty_request.resource_spans[0].scope_spans[0].spans[0]; + empty_span.events.clear(); + empty_span.links.clear(); + let result = instance + .traces( + instance.clone(), + empty_request, + PipelineWay::OtlpTraceDirectV2, + GreptimePipelineParams::default(), + "trace_v2".to_string(), + ctx.clone(), + ) + .await?; + assert_eq!((result.accepted_spans, result.rejected_spans), (1, 0)); + assert_trace_v2_query( + instance, + "SELECT COUNT(*) = 1 FROM trace_v2 + WHERE json_to_string(span_events) = '[]' AND json_to_string(span_links) = '[]'", + ctx.clone(), + ) + .await?; Ok(()) } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c6aa6d2f20d..12b9ac883d1 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -7827,11 +7827,21 @@ pub(crate) async fn test_otlp_traces_v2( "otlp_traces_v2_fixed_schema", &client, "select count(*) from information_schema.columns where table_name = 'trace_v2_spans' \ - and (column_name like 'span_attributes.%' \ - or column_name in ('span_events', 'span_links'));", + and column_name like 'span_attributes.%';", "[[0]]", ) .await; + validate_data( + "otlp_traces_v2_events_links", + &client, + r#"select json_get_string(span_events, '$[0].name'), + json_get_int(span_events, '$[0].attributes."event.code"'), + json_get_string(span_links, '$[0].trace_id'), + json_get_string(span_links, '$[0].attributes."link.type"') + from trace_v2_spans;"#, + r#"[["cache.hit",7,"cc9e0991a2e63d274984bd44ee669203","follows_from"]]"#, + ) + .await; validate_data( "otlp_traces_v2_escaped_keys", &client, diff --git a/tests/compatibility/cases/trace_v2_json2/case.toml b/tests/compatibility/cases/trace_v2_json2/case.toml index f4136b4e2d8..0f3f6f202bf 100644 --- a/tests/compatibility/cases/trace_v2_json2/case.toml +++ b/tests/compatibility/cases/trace_v2_json2/case.toml @@ -1,5 +1,5 @@ name = "trace_v2_json2" -reason = "Preserve the trace v2 schema, JSON2 attributes and model metadata across restart and upgrade." +reason = "Preserve the trace v2 schema, JSON2 attributes, JSON events/links and model metadata across restart and upgrade." introduced_by = "greptime_trace_v2" topologies = ["standalone", "distributed"] from_range = [">=v1.3.0-alpha.1"] diff --git a/tests/compatibility/cases/trace_v2_json2/setup.sql b/tests/compatibility/cases/trace_v2_json2/setup.sql index 02538b0ea70..70feb0ebaba 100644 --- a/tests/compatibility/cases/trace_v2_json2/setup.sql +++ b/tests/compatibility/cases/trace_v2_json2/setup.sql @@ -15,7 +15,9 @@ CREATE TABLE t_trace_v2 ( service_name STRING PRIMARY KEY, span_attributes JSON2, scope_attributes JSON2, - resource_attributes JSON2 + resource_attributes JSON2, + span_events JSON, + span_links JSON ) WITH ( 'append_mode' = 'true', 'table_data_model' = 'greptime_trace_v2', @@ -26,8 +28,16 @@ CREATE TABLE t_trace_v2 ( 'greptime.semantic.trace.conventions' = 'unknown' ); -INSERT INTO t_trace_v2 ("timestamp", timestamp_end, duration_nano, service_name, trace_id, span_id, span_attributes, scope_attributes, resource_attributes) VALUES - (1, 11, 10, NULL, 'trace-1', 'span-1', '{"http.status_code":200,"nested":{"a.b":[true,"ok"]}}', '{}', '{}'), - (1, 21, 20, NULL, 'trace-2', 'span-2', '{"http.status_code":"ok","bytes":[1,2,3]}', '{"enabled":true}', '{"deployment.environment":"prod"}'); +INSERT INTO t_trace_v2 ( + "timestamp", timestamp_end, duration_nano, service_name, trace_id, span_id, + span_attributes, scope_attributes, resource_attributes, span_events, span_links +) VALUES + (1, 11, 10, NULL, 'trace-1', 'span-1', + '{"http.status_code":200,"nested":{"a.b":[true,"ok"]}}', '{}', '{}', + '[{"name":"cache.hit","time":"2025-01-10 03:49:02.444400+0000","attributes":{"event.code":7}}]', + '[{"trace_id":"linked-trace","span_id":"linked-span","trace_state":"vendor=value","attributes":{"link.type":"follows_from"}}]'), + (1, 21, 20, NULL, 'trace-2', 'span-2', + '{"http.status_code":"ok","bytes":[1,2,3]}', + '{"enabled":true}', '{"deployment.environment":"prod"}', '[]', '[]'); ADMIN FLUSH_TABLE('t_trace_v2'); diff --git a/tests/compatibility/cases/trace_v2_json2/verify.result b/tests/compatibility/cases/trace_v2_json2/verify.result index 7d3d9f2bf7f..ebaaefad1d9 100644 --- a/tests/compatibility/cases/trace_v2_json2/verify.result +++ b/tests/compatibility/cases/trace_v2_json2/verify.result @@ -34,7 +34,7 @@ WHERE table_name = 't_trace_v2'; +----+ | n | +----+ -| 17 | +| 19 | +----+ SELECT COUNT(*) AS n FROM information_schema.tables @@ -47,3 +47,24 @@ WHERE table_name = 't_trace_v2' +---+ | 1 | +---+ + +SELECT COUNT(*) AS n FROM t_trace_v2 +WHERE json_get_string(span_events, '$[0].name') = 'cache.hit' + AND json_get_int(span_events, '$[0].attributes."event.code"') = 7 + AND json_get_string(span_links, '$[0].trace_id') = 'linked-trace' + AND json_get_string(span_links, '$[0].attributes."link.type"') = 'follows_from'; + ++---+ +| n | ++---+ +| 1 | ++---+ + +SELECT COUNT(*) AS n FROM t_trace_v2 +WHERE json_to_string(span_events) = '[]' AND json_to_string(span_links) = '[]'; + ++---+ +| n | ++---+ +| 1 | ++---+ diff --git a/tests/compatibility/cases/trace_v2_json2/verify.sql b/tests/compatibility/cases/trace_v2_json2/verify.sql index 7a1b0cd50bf..c6223bd7395 100644 --- a/tests/compatibility/cases/trace_v2_json2/verify.sql +++ b/tests/compatibility/cases/trace_v2_json2/verify.sql @@ -17,3 +17,12 @@ SELECT COUNT(*) AS n FROM information_schema.tables WHERE table_name = 't_trace_v2' AND create_options LIKE '%table_data_model=greptime_trace_v2%' AND create_options LIKE '%greptime.semantic.pipeline=greptime_trace_v2%'; + +SELECT COUNT(*) AS n FROM t_trace_v2 +WHERE json_get_string(span_events, '$[0].name') = 'cache.hit' + AND json_get_int(span_events, '$[0].attributes."event.code"') = 7 + AND json_get_string(span_links, '$[0].trace_id') = 'linked-trace' + AND json_get_string(span_links, '$[0].attributes."link.type"') = 'follows_from'; + +SELECT COUNT(*) AS n FROM t_trace_v2 +WHERE json_to_string(span_events) = '[]' AND json_to_string(span_links) = '[]'; From 27090496d1f975543020b5daf2d0e4b26fa679e4 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 18 Sep 2026 08:03:07 +0000 Subject: [PATCH 07/63] test(flow): stabilize FLUSH_FLOW assertions after async source-table mirrors (#9230) * test(flow): wait for async source-table mirrors before FLUSH_FLOW Streaming flows execute inline in the flownode insert handler since #8976, and source-table inserts are mirrored from the frontend as detached tasks. Tests that INSERT then ADMIN FLUSH_FLOW then SELECT intermittently miss rows on slow CI runners because the flush no longer implies the mirror reached the flownode (flow_no_aggr lost row 'l', flow_advance_ttl lost row '23' after the post-restart reinsert). Add SQLNESS SLEEP 3s between those mirror inserts and the flush, the established pattern in the flow suite, so the assertions are stable. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): also guard first-section flushes in flow_advance_ttl Local reproduction without sleeps failed in a window the previous commit did not cover: the first INSERT (20,20,22) of each section is followed immediately by ADMIN FLUSH_FLOW, and the distributed run observed an empty sink on that SELECT (~1 in 8 iterations). Add the same SLEEP 3s guard before those two flushes. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../cases/standalone/common/flow/flow_advance_ttl.result | 8 ++++++++ tests/cases/standalone/common/flow/flow_advance_ttl.sql | 8 ++++++++ tests/cases/standalone/common/flow/flow_no_aggr.result | 4 ++++ tests/cases/standalone/common/flow/flow_no_aggr.sql | 4 ++++ 4 files changed, 24 insertions(+) diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.result b/tests/cases/standalone/common/flow/flow_advance_ttl.result index 6d87854b961..279a6552d21 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.result +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.result @@ -86,6 +86,8 @@ VALUES Affected Rows: 0 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -146,6 +148,8 @@ VALUES Affected Rows: 0 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -283,6 +287,8 @@ VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -366,6 +372,8 @@ VALUES Affected Rows: 1 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.sql b/tests/cases/standalone/common/flow/flow_advance_ttl.sql index 9d7e12bf760..653fa16637f 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.sql +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.sql @@ -39,6 +39,8 @@ VALUES (20, "2021-07-01 00:00:00.200"), (22, "2021-07-01 00:00:00.600"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -65,6 +67,8 @@ VALUES (23, "2021-07-01 00:00:01.600"), (23, "2021-07-01 00:00:01.600"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -138,6 +142,8 @@ VALUES (20, "2021-07-01 00:00:00.200"), (22, "2021-07-01 00:00:00.600"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); @@ -160,6 +166,8 @@ INSERT INTO VALUES (23, "2021-07-01 00:00:01.600"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_distinct_basic'); diff --git a/tests/cases/standalone/common/flow/flow_no_aggr.result b/tests/cases/standalone/common/flow/flow_no_aggr.result index 94301ffb7f3..7ae7846e868 100644 --- a/tests/cases/standalone/common/flow/flow_no_aggr.result +++ b/tests/cases/standalone/common/flow/flow_no_aggr.result @@ -35,6 +35,8 @@ INSERT INTO base VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_out'); @@ -102,6 +104,8 @@ INSERT INTO base VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_out'); diff --git a/tests/cases/standalone/common/flow/flow_no_aggr.sql b/tests/cases/standalone/common/flow/flow_no_aggr.sql index bb0ed65f787..b0016dacf9b 100644 --- a/tests/cases/standalone/common/flow/flow_no_aggr.sql +++ b/tests/cases/standalone/common/flow/flow_no_aggr.sql @@ -21,6 +21,8 @@ INSERT INTO base VALUES ('j', '2023-01-01 00:00:09'), ('l', '2023-01-01 00:00:08'); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_out'); @@ -53,6 +55,8 @@ INSERT INTO base VALUES ('j', '2023-01-01 00:00:09'), ('l', '2023-01-01 00:00:08'); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_out'); From cb6c7a5c6d326ab37f1c7247ffc1ebaf0f57719e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 18 Sep 2026 09:51:01 +0000 Subject: [PATCH 08/63] perf(mito-codec): streamline sparse label extraction and buffer sizing (#9217) * perf(mito-codec): validate sparse labels once Signed-off-by: Lei, HUANG * perf(mito-codec): reserve complete sparse label chunks Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito-codec/src/index.rs | 30 +------ src/mito-codec/src/row_converter/sparse.rs | 22 +++++- .../src/row_converter/sparse/checked.rs | 79 +++++++++++++++++++ src/mito2/benches/bench_pk_tag_column.rs | 22 +++++- src/mito2/src/series_index/writer.rs | 22 +----- src/mito2/src/sst/parquet/flat_format.rs | 18 +---- 6 files changed, 126 insertions(+), 67 deletions(-) diff --git a/src/mito-codec/src/index.rs b/src/mito-codec/src/index.rs index 53ff5cd3c04..968173b878a 100644 --- a/src/mito-codec/src/index.rs +++ b/src/mito-codec/src/index.rs @@ -24,9 +24,7 @@ use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::ColumnMetadata; use store_api::storage::ColumnId; -use crate::error::{ - FieldTypeMismatchSnafu, IndexEncodeNullSnafu, InvalidSparsePrimaryKeySnafu, Result, -}; +use crate::error::{FieldTypeMismatchSnafu, IndexEncodeNullSnafu, Result}; use crate::row_converter::sparse::{ RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView, }; @@ -66,34 +64,14 @@ impl IndexValueCodec { column_id: ColumnId, buffer: &'a mut Vec, ) -> Result> { - let Some(encoded) = pk.encoded_value(column_id)? else { - return Ok(None); - }; - if encoded[0] == 0 { - return Ok(None); - } if matches!( column_id, RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID ) { - return Ok(Some(encoded)); + return pk.encoded_value(column_id); } - - buffer.clear(); - // The view has checked the Option marker, bytes marker and every chunk length. - // Reserving once avoids growing the buffer for each 8-byte chunk. - buffer.reserve(encoded.len() - 2); - for chunk in encoded[2..].chunks_exact(9) { - let len = usize::from(chunk[8]).min(8); - buffer.extend_from_slice(&chunk[..len]); - } - std::str::from_utf8(buffer).map_err(|_| { - InvalidSparsePrimaryKeySnafu { - reason: "label is not valid UTF-8", - } - .build() - })?; - Ok(Some(buffer.as_slice())) + pk.label(column_id, buffer) + .map(|value| value.map(str::as_bytes)) } /// Serializes a non-null `ValueRef` using the data type defined in `SortField` and writes diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index b2b3e0e3778..1be2fbacb2b 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -310,7 +310,7 @@ impl SparsePrimaryKeyCodec { { for (tag_column_id, tag_value) in row { let value_len = tag_value.len(); - buffer.reserve(6 + value_len / 8 * 9); + buffer.reserve(6 + value_len.div_ceil(8) * 9); buffer.put_u32(tag_column_id); buffer.put_u8(1); buffer.put_u8(!tag_value.is_empty() as u8); @@ -814,6 +814,26 @@ mod tests { assert_eq!(buffer, buffer_by_raw_encoding); } + #[test] + fn raw_encoding_matches_serde_at_chunk_boundaries() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut actual = Vec::new(); + for len in [0, 7, 8, 9, 15, 16, 17, 1024] { + let label = "x".repeat(len); + let mut expected = Vec::new(); + codec.encode_internal(42, 7, &mut expected).unwrap(); + codec + .encode_to_vec([(1, ValueRef::String(&label))].into_iter(), &mut expected) + .unwrap(); + actual.clear(); + codec.encode_internal(42, 7, &mut actual).unwrap(); + codec + .encode_raw_tag_value([(1, label.as_bytes())].into_iter(), &mut actual) + .unwrap(); + assert_eq!(actual, expected, "label length {len}"); + } + } + #[test] fn test_encode_to_vec() { let region_metadata = test_region_metadata(); diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs index 667f16fe804..1e57c13cc7f 100644 --- a/src/mito-codec/src/row_converter/sparse/checked.rs +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -86,6 +86,47 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { Ok(None) } + /// Extracts a UTF-8 label into a reusable buffer, validating it once. + /// Missing and null labels return `None` without modifying the buffer; + /// an empty label returns `Some("")`. Reserved numeric columns are rejected. + /// The returned string borrows `buffer`, whose contents are replaced for each label. + pub fn label<'buf>( + &mut self, + column_id: ColumnId, + buffer: &'buf mut Vec, + ) -> Result> { + ensure!( + !matches!( + column_id, + RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID + ), + InvalidSparsePrimaryKeySnafu { + reason: "reserved column is not a string label" + } + ); + let Some(encoded) = self.encoded_value(column_id)? else { + return Ok(None); + }; + if encoded[0] == 0 { + return Ok(None); + } + + buffer.clear(); + // encoded_value checked the markers and every chunk length. Reserve once + // before unchunking so the buffer does not grow for each 8-byte chunk. + buffer.reserve(encoded.len() - 2); + for chunk in encoded[2..].chunks_exact(9) { + let len = usize::from(chunk[8]).min(8); + buffer.extend_from_slice(&chunk[..len]); + } + std::str::from_utf8(buffer).map(Some).map_err(|_| { + InvalidSparsePrimaryKeySnafu { + reason: "label is not valid UTF-8", + } + .build() + }) + } + /// Returns the table id from the non-null prefix validated by [`Self::new`]. pub fn table_id(&self) -> u32 { (&self.pk[TABLE_ID_VALUE_OFFSET + 1..]).get_u32() @@ -265,4 +306,42 @@ mod tests { let mut view = SparsePrimaryKeyView::new(&invalid_utf8, &mut cache).unwrap(); assert!(IndexValueCodec::encode_sparse_value(&mut view, 1, &mut Vec::new()).is_err()); } + + #[test] + fn label_reuses_buffer_across_null_empty_invalid_and_valid_values() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut cache = SparseOffsetsCache::new(); + let mut buffer = Vec::with_capacity(256); + let ptr = buffer.as_ptr(); + for label in [ + b"1234567\xe4\xb8\xad\0\xe6\x96\x87".as_slice(), + b"", + b"\xff", + b"after-error", + ] { + let mut pk = Vec::new(); + codec.encode_internal(42, 7, &mut pk).unwrap(); + codec + .encode_raw_tag_value([(1, label)].into_iter(), &mut pk) + .unwrap(); + pk.extend_from_slice(&2_u32.to_be_bytes()); + pk.push(0); + let mut view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap(); + let actual = view.label(1, &mut buffer); + match std::str::from_utf8(label) { + Ok(expected) => assert_eq!(actual.unwrap(), Some(expected)), + Err(_) => assert!(actual.is_err()), + } + assert_eq!(view.label(2, &mut buffer).unwrap(), None); + assert_eq!(view.label(3, &mut buffer).unwrap(), None); + // Null/missing values do not overwrite the scratch buffer, matching index extraction. + assert_eq!(buffer, label); + assert_eq!(buffer.as_ptr(), ptr); + assert!( + view.label(RESERVED_COLUMN_ID_TABLE_ID, &mut buffer) + .is_err() + ); + assert!(view.label(RESERVED_COLUMN_ID_TSID, &mut buffer).is_err()); + } + } } diff --git a/src/mito2/benches/bench_pk_tag_column.rs b/src/mito2/benches/bench_pk_tag_column.rs index 7f10c912138..d3ded3d0b05 100644 --- a/src/mito2/benches/bench_pk_tag_column.rs +++ b/src/mito2/benches/bench_pk_tag_column.rs @@ -95,6 +95,10 @@ fn metadata(tags: u32) -> RegionMetadataRef { /// Builds a sparse flat batch whose primary key dictionary holds /// `ROWS / rows_per_key` distinct keys with `tags` labels each. fn input(tags: u32, rows_per_key: usize) -> RecordBatch { + input_with_label_len(tags, rows_per_key, 24) +} + +fn input_with_label_len(tags: u32, rows_per_key: usize, label_len: usize) -> RecordBatch { let codec = SparsePrimaryKeyCodec::schemaless(); let mut keys = BinaryDictionaryBuilder::::new(); for series in 0..ROWS / rows_per_key { @@ -103,7 +107,14 @@ fn input(tags: u32, rows_per_key: usize) -> RecordBatch { .encode_internal((series / 128) as u32, series as u64, &mut key) .unwrap(); let labels: Vec<_> = (0..tags) - .map(|id| (id, format!("tag-{id:03}-value-{series:010}"))) + .map(|id| { + let mut value = format!("tag-{id:03}-value-{series:010}"); + value.extend(std::iter::repeat_n( + 'x', + label_len.saturating_sub(value.len()), + )); + (id, value) + }) .collect(); codec .encode_raw_tag_value( @@ -203,8 +214,8 @@ fn bench_pk_tag_filters(c: &mut Criterion) { const TAGS: u32 = 40; let metadata = metadata(TAGS); let mut group = c.benchmark_group("pk_tag_filters"); - for rows_per_key in [1, 32] { - let batch = input(TAGS, rows_per_key); + for (rows_per_key, label_len) in [(1, 24), (32, 24), (1, 1024), (32, 1024)] { + let batch = input_with_label_len(TAGS, rows_per_key, label_len); for predicate_count in [1, 2, 4, 8, 16, 32] { // Use distinct tags at the end of the key to exercise offset discovery. // Every row matches, so increasing the predicate count does not change selectivity. @@ -215,7 +226,10 @@ fn bench_pk_tag_filters(c: &mut Criterion) { // Validate the workload outside the timed section. assert_eq!(filter(batch.clone()).unwrap().unwrap().num_rows(), ROWS); group.bench_function( - BenchmarkId::new(format!("{rows_per_key}rpk"), predicate_count), + BenchmarkId::new( + format!("{rows_per_key}rpk_{label_len}bytes"), + predicate_count, + ), |b| b.iter(|| black_box(filter(black_box(batch.clone())).unwrap())), ); } diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index b6330aac044..0863b3b8519 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -26,7 +26,6 @@ use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::timestamp::timestamp_array_to_primitive; -use mito_codec::index::IndexValueCodec; use mito_codec::row_converter::SparseOffsetsCache; use mito_codec::row_converter::sparse::SparsePrimaryKeyView; use object_store::ObjectStore; @@ -603,25 +602,8 @@ fn decode_primary_key( let mut tags = Vec::with_capacity(tag_columns.len()); for (column_id, _) in tag_columns { - // `encode_sparse_value` returns None for missing and null labels and - // validates UTF-8 for string labels. - let value = IndexValueCodec::encode_sparse_value(&mut view, *column_id, buf) - .context(DecodeSnafu)?; - let tag = match value { - None => None, - Some(bytes) => { - let value = std::str::from_utf8(bytes).map_err(|_| { - InvalidRecordBatchSnafu { - reason: format!( - "sparse tag value of column {column_id} is not valid UTF-8" - ), - } - .build() - })?; - Some(value.to_string()) - } - }; - tags.push(tag); + let value = view.label(*column_id, buf).context(DecodeSnafu)?; + tags.push(value.map(str::to_owned)); } Ok(SeriesIndexRow { diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index c692dae0908..a99fc201baf 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -42,7 +42,6 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::value::ValueRef; use datatypes::vectors::MutableVector; -use mito_codec::index::IndexValueCodec; use mito_codec::row_converter::sparse::{ RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView, }; @@ -798,23 +797,10 @@ fn push_sparse_tag_value_in_view( RESERVED_COLUMN_ID_TABLE_ID => builder.push_value_ref(&ValueRef::UInt32(view.table_id())), RESERVED_COLUMN_ID_TSID => builder.push_value_ref(&ValueRef::UInt64(view.tsid())), _ => { - // `encode_sparse_value` returns None for missing and null labels - // and validates UTF-8 for string labels. - let value = IndexValueCodec::encode_sparse_value(view, column_id, value_buf) - .context(DecodeSnafu)?; + let value = view.label(column_id, value_buf).context(DecodeSnafu)?; match value { None => builder.push_null(), - Some(bytes) => { - let value = std::str::from_utf8(bytes).map_err(|_| { - InvalidRecordBatchSnafu { - reason: format!( - "sparse tag value of column {column_id} is not valid UTF-8" - ), - } - .build() - })?; - builder.push_value_ref(&ValueRef::String(value)); - } + Some(value) => builder.push_value_ref(&ValueRef::String(value)), } } } From 983738101a369a8d3db476f6fa08975536acf061 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 18 Sep 2026 11:34:32 +0000 Subject: [PATCH 09/63] feat(flow): admit mergeable average states in incremental plans (#9235) Incremental batching flows already merge sink state for scalar aggregates and the HLL/UDDSketch/stddev state families. AVG now exposes a mergeable Binary state (avg_state / avg_merge with __avg_state_delta_merge), so admit those aggregates too instead of forcing a full snapshot for flows whose only state column is an average. merge_op_for_aggregate_expr takes the aggregate input schema so the avg_merge arm can require a Binary state argument; the state form is the aggregate result persisted by the sink, so no extra coercion is needed. Other input types keep being rejected. Also cover avg_state, avg_merge and duplicate AVG projections in the incremental plan analysis tests, extend the mixed state-family rewrite test with an average column, and extend the standalone partitioned state-merge SQLness case with avg_state/avg_merge compared against a direct avg over the source. Signed-off-by: discord9 Co-authored-by: discord9 --- src/flow/src/batching_mode/utils.rs | 20 ++- src/flow/src/batching_mode/utils/test.rs | 170 +++++++++++++++++- .../flow/flow_incremental_state_merge.result | 46 ++--- .../flow/flow_incremental_state_merge.sql | 16 +- 4 files changed, 223 insertions(+), 29 deletions(-) diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 87bb0925061..0220f4ef48d 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -21,6 +21,7 @@ use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_function::aggrs::aggr_wrapper::get_aggr_func; use common_telemetry::debug; +use datafusion::arrow::datatypes::DataType as ArrowDataType; use datafusion::datasource::DefaultTableSource; use datafusion::error::Result as DfResult; use datafusion::logical_expr::Expr; @@ -325,7 +326,10 @@ fn is_literal_or_cast_literal(expr: &Expr) -> bool { } } -fn merge_op_for_aggregate_expr(aggr_expr: &Expr) -> Result { +fn merge_op_for_aggregate_expr( + aggr_expr: &Expr, + input_schema: &DFSchema, +) -> Result { let Some(aggr_func) = get_aggr_func(aggr_expr) else { return Err(aggr_expr.to_string()); }; @@ -355,6 +359,9 @@ fn merge_op_for_aggregate_expr(aggr_expr: &Expr) -> Result Ok(IncrementalAggregateMergeOp::BitOr), "bit_xor" => Ok(IncrementalAggregateMergeOp::BitXor), // Preserve state-family parameters; value coercion is handled by the aggregate. + "avg_state" if aggr_func.params.args.len() == 1 => { + state_delta_merge("__avg_state_delta_merge", vec![]) + } "hll" if aggr_func.params.args.len() == 1 => state_delta_merge("__hll_delta_merge", vec![]), "stddev_pop_state" if aggr_func.params.args.len() == 1 => { state_delta_merge("__stddev_pop_state_delta_merge", vec![]) @@ -373,6 +380,15 @@ fn merge_op_for_aggregate_expr(aggr_expr: &Expr) -> Result + { + state_delta_merge("__avg_state_delta_merge", vec![]) + } _ => Err(aggr_expr.to_string()), } } @@ -519,7 +535,7 @@ pub fn analyze_incremental_aggregate_plan( &group_key_names, )); for aggr_expr in aggr_exprs { - let merge_op = match merge_op_for_aggregate_expr(&aggr_expr) { + let merge_op = match merge_op_for_aggregate_expr(&aggr_expr, aggregate.input.schema()) { Ok(merge_op) => merge_op, Err(reason) => { unsupported_exprs.push(reason); diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 7b3099818be..6315ae3d679 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -16,6 +16,7 @@ use std::collections::BTreeMap; use std::sync::Arc; use catalog::RegisterTableRequest; +use common_function::aggrs::approximate::avg::AvgState; use common_query::OutputData; use common_recordbatch::recordbatch::merge_record_batches; use common_recordbatch::{RecordBatch, util}; @@ -1706,11 +1707,154 @@ async fn test_analyze_incremental_aggregate_plan_rejects_avg() { assert!(!analysis.unsupported_exprs.is_empty()); } +#[tokio::test] +async fn test_analyze_incremental_aggregate_plan_supports_avg_state() { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let sql = "SELECT avg_state(number) AS avg_num, ts FROM numbers_with_ts GROUP BY ts"; + let plan = sql_to_df_plan(ctx, query_engine, sql, false).await.unwrap(); + + let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); + assert!( + analysis.unsupported_exprs.is_empty(), + "avg_state should be supported: {:?}", + analysis.unsupported_exprs + ); + assert_eq!(analysis.merge_columns.len(), 1); + assert_eq!(analysis.merge_columns[0].output_field_name, "avg_num"); + assert!(matches!( + &analysis.merge_columns[0].merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + )); +} + +#[tokio::test] +async fn test_analyze_incremental_aggregate_plan_supports_avg_merge() { + let query_engine = create_test_query_engine(); + let producer_sql = "SELECT avg_state(number) AS avg_state, CASE WHEN number <= 5 THEN 1 ELSE 2 END AS grp FROM numbers_with_ts GROUP BY grp"; + let producer_plan = sql_to_df_plan( + QueryContext::arc(), + query_engine.clone(), + producer_sql, + false, + ) + .await + .unwrap(); + let producer_output = query_engine + .execute(producer_plan, QueryContext::arc()) + .await + .unwrap(); + let OutputData::Stream(producer_stream) = producer_output.data else { + panic!("expected AVG state producer execution to be a stream"); + }; + let producer_batches = util::collect(producer_stream).await.unwrap(); + let producer_schema = producer_batches.first().unwrap().schema.clone(); + let avg_states = merge_record_batches(producer_schema, &producer_batches).unwrap(); + assert_eq!(avg_states.num_rows(), 2); + let avg_states_table = MemTable::table("avg_states", avg_states); + query_engine + .engine_state() + .catalog_manager() + .as_any() + .downcast_ref::() + .unwrap() + .register_table_sync(RegisterTableRequest { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table_name: "avg_states".to_string(), + table_id: 4096, + table: avg_states_table, + }) + .unwrap(); + + let sql = "SELECT avg_merge(avg_state) AS avg_num, grp FROM avg_states GROUP BY grp"; + let plan = sql_to_df_plan(QueryContext::arc(), query_engine.clone(), sql, false) + .await + .unwrap(); + let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); + assert!( + analysis.unsupported_exprs.is_empty(), + "avg_merge should be supported: {:?}", + analysis.unsupported_exprs + ); + assert_eq!(analysis.merge_columns.len(), 1); + assert!(matches!( + &analysis.merge_columns[0].merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + )); + + let output = query_engine + .execute(plan, QueryContext::arc()) + .await + .unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected AVG merge execution to be a stream"); + }; + let batches = util::collect(stream).await.unwrap(); + let schema = batches.first().unwrap().schema.clone(); + let batch = merge_record_batches(schema, &batches).unwrap(); + let groups = batch + .column_by_name("grp") + .unwrap() + .as_primitive::(); + let averages = batch.column_by_name("avg_num").unwrap().as_binary::(); + let mut values = (0..batch.num_rows()) + .map(|index| { + ( + groups.value(index), + AvgState::decode(averages.value(index)).unwrap().average(), + ) + }) + .collect::>(); + values.sort_unstable_by_key(|(group, _)| *group); + assert_eq!(values, [(1, Some(3.0)), (2, Some(8.0))]); +} + +#[tokio::test] +async fn test_analyze_incremental_aggregate_plan_supports_duplicate_avg_projections() { + let analysis = analyze_test_sql( + "SELECT avg_state(number) AS avg_a, avg_state(number) AS avg_b, avg_state(number + 1) AS avg_num_plus, ts FROM numbers_with_ts GROUP BY ts", + ) + .await; + + assert!(analysis.unsupported_exprs.is_empty()); + assert_eq!(analysis.merge_columns.len(), 3); + assert!(analysis.merge_columns.iter().all(|column| { + matches!( + &column.merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + ) + })); + assert_eq!(analysis.merge_columns[0].output_field_name, "avg_a"); + assert_eq!(analysis.merge_columns[0].input_field_name, "avg_a"); + let avg_b = analysis + .merge_columns + .iter() + .find(|column| column.output_field_name == "avg_b") + .unwrap(); + assert_eq!(avg_b.input_field_name, "avg_a"); + assert!( + analysis + .merge_columns + .iter() + .any(|column| column.output_field_name == "avg_num_plus") + ); +} + #[tokio::test] async fn test_rewrite_incremental_aggregate_merges_populated_mixed_state_families() { let query_engine = create_test_query_engine(); - let old_sql = "SELECT hll(CAST(number AS VARCHAR)) AS hll_a, hll(CAST(number AS VARCHAR)) AS hll_b, uddsketch_state(128, 0.000001, number) AS percentile_a, uddsketch_state(256, 0.02, number) AS percentile_b, stddev_pop_state(number) AS stddev_state, sum(number) AS total, CASE WHEN number <= 5 THEN CAST(NULL AS BIGINT) ELSE 2 END AS grp FROM numbers_with_ts WHERE number <= 3 OR number = 6 GROUP BY grp"; - let new_sql = "SELECT hll(CAST(number AS VARCHAR)) AS hll_a, hll(CAST(number AS VARCHAR)) AS hll_b, uddsketch_state(128, 0.000001, number) AS percentile_a, uddsketch_state(256, 0.02, number) AS percentile_b, stddev_pop_state(number) AS stddev_state, sum(number) AS total, CASE WHEN number <= 5 THEN CAST(NULL AS BIGINT) WHEN number <= 8 THEN 2 ELSE 3 END AS grp FROM numbers_with_ts WHERE number >= 4 AND number != 6 GROUP BY grp"; + let old_sql = "SELECT avg_state(number) AS avg_num, hll(CAST(number AS VARCHAR)) AS hll_a, hll(CAST(number AS VARCHAR)) AS hll_b, uddsketch_state(128, 0.000001, number) AS percentile_a, uddsketch_state(256, 0.02, number) AS percentile_b, stddev_pop_state(number) AS stddev_state, sum(number) AS total, CASE WHEN number <= 5 THEN CAST(NULL AS BIGINT) ELSE 2 END AS grp FROM numbers_with_ts WHERE number <= 3 OR number = 6 GROUP BY grp"; + let new_sql = "SELECT avg_state(number) AS avg_num, hll(CAST(number AS VARCHAR)) AS hll_a, hll(CAST(number AS VARCHAR)) AS hll_b, uddsketch_state(128, 0.000001, number) AS percentile_a, uddsketch_state(256, 0.02, number) AS percentile_b, stddev_pop_state(number) AS stddev_state, sum(number) AS total, CASE WHEN number <= 5 THEN CAST(NULL AS BIGINT) WHEN number <= 8 THEN 2 ELSE 3 END AS grp FROM numbers_with_ts WHERE number >= 4 AND number != 6 GROUP BY grp"; let old_plan = sql_to_df_plan(QueryContext::arc(), query_engine.clone(), old_sql, false) .await .unwrap(); @@ -1757,6 +1901,7 @@ async fn test_rewrite_incremental_aggregate_merges_populated_mixed_state_familie assert_eq!( analysis.output_field_names, vec![ + "avg_num", "hll_a", "hll_b", "percentile_a", @@ -1783,6 +1928,27 @@ async fn test_rewrite_incremental_aggregate_merges_populated_mixed_state_familie 3, "rewrite must produce one row per group" ); + let avg_states = merged_batch + .column_by_name("avg_num") + .unwrap() + .as_binary::(); + let merged_groups = merged_batch + .column_by_name("grp") + .unwrap() + .as_primitive::(); + assert_eq!(avg_states.null_count(), 0); + for index in 0..merged_batch.num_rows() { + let expected = match (!merged_groups.is_null(index)).then(|| merged_groups.value(index)) { + None => Some(3.0), + Some(2) => Some(7.0), + Some(3) => Some(9.5), + group => panic!("unexpected group: {group:?}"), + }; + assert_eq!( + AvgState::decode(avg_states.value(index)).unwrap().average(), + expected + ); + } let merged_table = MemTable::table("merged_states", merged_batch); query_engine .engine_state() diff --git a/tests/cases/standalone/common/flow/flow_incremental_state_merge.result b/tests/cases/standalone/common/flow/flow_incremental_state_merge.result index 9a9b23dfef3..49b112f564b 100644 --- a/tests/cases/standalone/common/flow/flow_incremental_state_merge.result +++ b/tests/cases/standalone/common/flow/flow_incremental_state_merge.result @@ -1,7 +1,7 @@ -- Incremental aggregate state merge over a partitioned source (two regions). -- A WHERE-filtered flow mixes scalar aggregates (sum, count) with aggregate --- state columns (hll, stddev_pop_state, uddsketch_state); each later delta in --- the same time window must merge its state with the state already persisted +-- state columns (hll, stddev_pop_state, uddsketch_state, avg_state); each later +-- delta in the same time window must merge its state with the state already persisted -- in the sink. Covers two nonempty deltas across both regions, one round whose -- delta is empty because the flow's predicate drops the newly inserted row, -- and a final nonempty delta checked against a direct source aggregation. @@ -32,6 +32,7 @@ SELECT hll(user_id) AS user_hll, stddev_pop_state(v) AS v_stddev_state, uddsketch_state(128, 0.01, v) AS v_sketch_state, + avg_state(v) AS v_avg_state, date_bin(INTERVAL '1 minute', ts, '2024-01-01 00:00:00') AS time_window FROM flow_incr_state_merge_input @@ -115,6 +116,7 @@ SELECT hll_count(hll_merge(user_hll)) AS users, ROUND(stddev_pop_calc(stddev_pop_merge(v_stddev_state)), 6) AS v_stddev, ROUND(uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)), 4) AS v_p50, + ROUND(avg_calc(avg_merge(v_avg_state)), 6) AS v_avg, time_window FROM flow_incr_state_merge_sink @@ -125,11 +127,11 @@ GROUP BY ORDER BY time_window; -+-------+-----------+-------+----------+--------+---------------------+ -| total | row_count | users | v_stddev | v_p50 | time_window | -+-------+-----------+-------+----------+--------+---------------------+ -| 10 | 4 | 4 | 1.118034 | 2.9742 | 2024-01-01T00:00:00 | -+-------+-----------+-------+----------+--------+---------------------+ ++-------+-----------+-------+----------+--------+-------+---------------------+ +| total | row_count | users | v_stddev | v_p50 | v_avg | time_window | ++-------+-----------+-------+----------+--------+-------+---------------------+ +| 10 | 4 | 4 | 1.118034 | 2.9742 | 2.5 | 2024-01-01T00:00:00 | ++-------+-----------+-------+----------+--------+-------+---------------------+ ADMIN FLUSH_TABLE('flow_incr_state_merge_input'); @@ -199,6 +201,7 @@ SELECT hll_count(hll_merge(user_hll)) AS users, ROUND(stddev_pop_calc(stddev_pop_merge(v_stddev_state)), 6) AS v_stddev, ROUND(uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)), 4) AS v_p50, + ROUND(avg_calc(avg_merge(v_avg_state)), 6) AS v_avg, time_window FROM flow_incr_state_merge_sink @@ -209,11 +212,11 @@ GROUP BY ORDER BY time_window; -+-------+-----------+-------+----------+--------+---------------------+ -| total | row_count | users | v_stddev | v_p50 | time_window | -+-------+-----------+-------+----------+--------+---------------------+ -| 10 | 4 | 4 | 1.118034 | 2.9742 | 2024-01-01T00:00:00 | -+-------+-----------+-------+----------+--------+---------------------+ ++-------+-----------+-------+----------+--------+-------+---------------------+ +| total | row_count | users | v_stddev | v_p50 | v_avg | time_window | ++-------+-----------+-------+----------+--------+-------+---------------------+ +| 10 | 4 | 4 | 1.118034 | 2.9742 | 2.5 | 2024-01-01T00:00:00 | ++-------+-----------+-------+----------+--------+-------+---------------------+ -- last_execution_time must have advanced past the mark taken above, proving -- the dirty window triggered a new round even though its delta was empty. A @@ -295,7 +298,8 @@ SELECT merged.row_count = direct.row_count AS count_matches, merged.users = direct.users AS hll_matches, abs(merged.v_stddev - direct.v_stddev) < 1e-9 AS stddev_matches, - ROUND(merged.v_p50, 4) = ROUND(direct.v_p50, 4) AS p50_matches + ROUND(merged.v_p50, 4) = ROUND(direct.v_p50, 4) AS p50_matches, + abs(merged.v_avg - direct.v_avg) < 1e-9 AS avg_matches FROM ( SELECT @@ -303,7 +307,8 @@ FROM row_count, hll_count(hll_merge(user_hll)) AS users, stddev_pop_calc(stddev_pop_merge(v_stddev_state)) AS v_stddev, - uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)) AS v_p50 + uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)) AS v_p50, + avg_calc(avg_merge(v_avg_state)) AS v_avg FROM flow_incr_state_merge_sink GROUP BY @@ -316,18 +321,19 @@ FROM count(n) AS row_count, hll_count(hll(user_id)) AS users, stddev_pop(v) AS v_stddev, - uddsketch_calc(0.5, uddsketch_state(128, 0.01, v)) AS v_p50 + uddsketch_calc(0.5, uddsketch_state(128, 0.01, v)) AS v_p50, + avg(v) AS v_avg FROM flow_incr_state_merge_input WHERE v > 0 ) AS direct; -+---------------+---------------+-------------+----------------+-------------+ -| total_matches | count_matches | hll_matches | stddev_matches | p50_matches | -+---------------+---------------+-------------+----------------+-------------+ -| true | true | true | true | true | -+---------------+---------------+-------------+----------------+-------------+ ++---------------+---------------+-------------+----------------+-------------+-------------+ +| total_matches | count_matches | hll_matches | stddev_matches | p50_matches | avg_matches | ++---------------+---------------+-------------+----------------+-------------+-------------+ +| true | true | true | true | true | true | ++---------------+---------------+-------------+----------------+-------------+-------------+ -- Clean up DROP FLOW flow_incr_state_merge; diff --git a/tests/cases/standalone/common/flow/flow_incremental_state_merge.sql b/tests/cases/standalone/common/flow/flow_incremental_state_merge.sql index a48b39e9fd2..595390f755e 100644 --- a/tests/cases/standalone/common/flow/flow_incremental_state_merge.sql +++ b/tests/cases/standalone/common/flow/flow_incremental_state_merge.sql @@ -1,7 +1,7 @@ -- Incremental aggregate state merge over a partitioned source (two regions). -- A WHERE-filtered flow mixes scalar aggregates (sum, count) with aggregate --- state columns (hll, stddev_pop_state, uddsketch_state); each later delta in --- the same time window must merge its state with the state already persisted +-- state columns (hll, stddev_pop_state, uddsketch_state, avg_state); each later +-- delta in the same time window must merge its state with the state already persisted -- in the sink. Covers two nonempty deltas across both regions, one round whose -- delta is empty because the flow's predicate drops the newly inserted row, -- and a final nonempty delta checked against a direct source aggregation. @@ -30,6 +30,7 @@ SELECT hll(user_id) AS user_hll, stddev_pop_state(v) AS v_stddev_state, uddsketch_state(128, 0.01, v) AS v_sketch_state, + avg_state(v) AS v_avg_state, date_bin(INTERVAL '1 minute', ts, '2024-01-01 00:00:00') AS time_window FROM flow_incr_state_merge_input @@ -77,6 +78,7 @@ SELECT hll_count(hll_merge(user_hll)) AS users, ROUND(stddev_pop_calc(stddev_pop_merge(v_stddev_state)), 6) AS v_stddev, ROUND(uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)), 4) AS v_p50, + ROUND(avg_calc(avg_merge(v_avg_state)), 6) AS v_avg, time_window FROM flow_incr_state_merge_sink @@ -131,6 +133,7 @@ SELECT hll_count(hll_merge(user_hll)) AS users, ROUND(stddev_pop_calc(stddev_pop_merge(v_stddev_state)), 6) AS v_stddev, ROUND(uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)), 4) AS v_p50, + ROUND(avg_calc(avg_merge(v_avg_state)), 6) AS v_avg, time_window FROM flow_incr_state_merge_sink @@ -194,7 +197,8 @@ SELECT merged.row_count = direct.row_count AS count_matches, merged.users = direct.users AS hll_matches, abs(merged.v_stddev - direct.v_stddev) < 1e-9 AS stddev_matches, - ROUND(merged.v_p50, 4) = ROUND(direct.v_p50, 4) AS p50_matches + ROUND(merged.v_p50, 4) = ROUND(direct.v_p50, 4) AS p50_matches, + abs(merged.v_avg - direct.v_avg) < 1e-9 AS avg_matches FROM ( SELECT @@ -202,7 +206,8 @@ FROM row_count, hll_count(hll_merge(user_hll)) AS users, stddev_pop_calc(stddev_pop_merge(v_stddev_state)) AS v_stddev, - uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)) AS v_p50 + uddsketch_calc(0.5, uddsketch_merge(128, 0.01, v_sketch_state)) AS v_p50, + avg_calc(avg_merge(v_avg_state)) AS v_avg FROM flow_incr_state_merge_sink GROUP BY @@ -215,7 +220,8 @@ FROM count(n) AS row_count, hll_count(hll(user_id)) AS users, stddev_pop(v) AS v_stddev, - uddsketch_calc(0.5, uddsketch_state(128, 0.01, v)) AS v_p50 + uddsketch_calc(0.5, uddsketch_state(128, 0.01, v)) AS v_p50, + avg(v) AS v_avg FROM flow_incr_state_merge_input WHERE From 5ccc29ce771c88de78929cb4e11973684b438799 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Sat, 19 Sep 2026 02:12:59 +0000 Subject: [PATCH 10/63] feat(log-store): add object store WAL store construction and recovery (#9238) * feat(log-store): add object store WAL store construction and recovery Signed-off-by: jeremyhi * fix(log-store): derive WAL store prefix from node and generation Signed-off-by: jeremyhi * docs(log-store): describe unsupported operation boundary Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/error.rs | 43 + src/log-store/src/object_store_wal.rs | 6 +- src/log-store/src/object_store_wal/store.rs | 1406 +++++++++++++++++++ 3 files changed, 1454 insertions(+), 1 deletion(-) create mode 100644 src/log-store/src/object_store_wal/store.rs diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 54fcc84b690..6826be716f3 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::any::Any; +use std::sync::Arc; use common_error::ext::{ErrorExt, RetryHint, retry_hint_from_io_error}; use common_error::status_code::StatusCode; @@ -370,6 +371,42 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display("Invalid WAL object, path: {}", path))] + InvalidWalObject { + path: String, + #[snafu(source(from(Error, Box::new)))] + source: Box, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Object store WAL prefix mismatch, expected: {}, actual: {}", + expected, + actual + ))] + MismatchedWalPrefix { + expected: String, + actual: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Object store WAL operation failed"))] + ObjectStoreWal { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + + /// The store supports construction, recovery, stopping and latest-entry queries; + /// other log-store operations return this error. + #[snafu(display("Object store WAL operation is not supported yet"))] + UnsupportedObjectStoreWalOperation { + #[snafu(implicit)] + location: Location, + }, } pub type Result = std::result::Result; @@ -402,6 +439,7 @@ impl ErrorExt for Error { | MissingValue { .. } | OverrideCompactedEntry { .. } | InvalidWalObjectStore { .. } + | MismatchedWalPrefix { .. } | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, StartWalTask { .. } | StopWalTask { .. } @@ -422,6 +460,10 @@ impl ErrorExt for Error { | WalObjectSequenceExhausted { .. } | WalEntryPositionExhausted { .. } => StatusCode::Unexpected, + UnsupportedObjectStoreWalOperation { .. } => StatusCode::Unsupported, + InvalidWalObject { source, .. } => source.status_code(), + ObjectStoreWal { source, .. } => source.status_code(), + // Object store related errors CreateWriter { .. } | WriteIndex { .. } @@ -457,6 +499,7 @@ impl ErrorExt for Error { | WriteIndex { error, .. } | ReadIndex { error, .. } | WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), + ObjectStoreWal { source, .. } => source.retry_hint(), Io { error, .. } => retry_hint_from_io_error(error), FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 98895869b81..70bc39dedca 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -43,7 +43,6 @@ //! id name the object that holds the entry, the low bits its position among //! the entries of its region in that object. -// These modules have no callers until the store that writes and reads objects lands. #[allow(dead_code)] mod batch; #[allow(dead_code)] @@ -53,5 +52,10 @@ mod format; #[allow(dead_code)] mod io; +#[allow(dead_code)] +mod store; + #[allow(unused_imports)] pub(crate) use batch::entry_id; +#[allow(unused_imports)] +pub(crate) use store::ObjectStoreLogStore; diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs new file mode 100644 index 00000000000..e578bc79126 --- /dev/null +++ b/src/log-store/src/object_store_wal/store.rs @@ -0,0 +1,1406 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Construction and footer-based recovery of an object store WAL. + +use std::collections::HashMap; +use std::fmt; +use std::ops::Range; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, PoisonError, RwLock}; +use std::time::Duration; + +use bytes::Bytes; +use common_wal::config::object_store::ObjectStoreWalConfig; +use futures::{StreamExt, TryStreamExt}; +use object_store::ObjectStore; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use store_api::logstore::entry::Entry; +use store_api::logstore::provider::{ObjectStoreProvider, Provider}; +use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; +use store_api::storage::RegionId; +use tokio::sync::{mpsc, oneshot}; + +use crate::error::{ + CorruptedWalObjectSnafu, Error, InvalidProviderSnafu, InvalidWalObjectSnafu, + InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, ObjectStoreWalSnafu, Result, + UnsupportedObjectStoreWalOperationSnafu, +}; +use crate::object_store_wal::catalog::ObjectCatalog; +use crate::object_store_wal::format::{ + FixedTrailer, FooterEntry, HEADER_LEN, MIN_OBJECT_LEN, TRAILER_LEN, decode_footer, + decode_header, decode_trailer, footer_range, verify_segment_ranges, +}; +use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; + +const COMMAND_BUFFER: usize = 1024; +const MIN_FLUSH_INTERVAL: Duration = Duration::from_millis(10); +/// Number of objects whose footers recovery fetches at a time. +const RECOVERY_CONCURRENCY: usize = 8; +/// Bytes recovery reads from the end of an object in one request. The window +/// holds the trailer and the footer of an object with up to 1364 regions of +/// 48 bytes each, so a second request for the footer is rare. +const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; + +/// A log store over the immutable WAL objects under one prefix. +pub(crate) struct ObjectStoreLogStore { + prefix: String, + io: Arc, + catalog: Arc>, + terminal_error: TerminalError, + stopped: Arc, + command_tx: mpsc::Sender, +} + +type TerminalError = Arc>>>; + +impl fmt::Debug for ObjectStoreLogStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ObjectStoreLogStore") + .field("prefix", &self.prefix) + .finish_non_exhaustive() + } +} + +impl ObjectStoreLogStore { + /// Builds the store under the node and generation prefix derived from `config`, + /// recovering the catalog from the objects that already exist. Recovery + /// fails on the first corrupted or conflicting object. + pub(crate) async fn try_new( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + node_id: u64, + generation: u64, + ) -> Result> { + let prefix = config.node_prefix(node_id, generation); + let io = ObjectStoreIo::new(object_store, &prefix)?; + Self::open(Arc::new(io), config, prefix).await + } + + async fn open( + io: Arc, + config: &ObjectStoreWalConfig, + prefix: String, + ) -> Result> { + ensure!( + config.flush_interval >= MIN_FLUSH_INTERVAL, + InvalidWalObjectStoreSnafu { + reason: format!( + "flush interval {:?} is shorter than {MIN_FLUSH_INTERVAL:?}", + config.flush_interval + ), + } + ); + + positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; + let (catalog, next_object_seq, durable_entry_ids) = recover(io.as_ref()).await?; + let catalog = Arc::new(RwLock::new(catalog)); + let terminal_error = TerminalError::default(); + let stopped = Arc::new(AtomicBool::new(false)); + let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + let actor = Actor { + terminal_error: terminal_error.clone(), + stopped: stopped.clone(), + command_rx, + issued_entry_ids: durable_entry_ids, + next_object_seq: Some(next_object_seq), + stop: Vec::new(), + }; + common_runtime::spawn_global(actor.run()); + Ok(Arc::new(Self { + prefix, + io, + catalog, + terminal_error, + stopped, + command_tx, + })) + } + + fn durable_entry_id_of(&self, region_id: RegionId) -> EntryId { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + catalog.region_max_entry_id(region_id).unwrap_or(0) + } + + /// Returns the region of `provider`, which must select this store's prefix. + fn region_of(&self, provider: &Provider) -> Result { + let provider = + provider + .as_object_store_provider() + .with_context(|| InvalidProviderSnafu { + expected: ObjectStoreProvider::type_name(), + actual: provider.type_name(), + })?; + ensure!( + provider.prefix == self.prefix, + MismatchedWalPrefixSnafu { + expected: self.prefix.clone(), + actual: provider.prefix.clone(), + } + ); + Ok(provider.region_id) + } + + fn check_terminal(&self) -> Result<()> { + match terminal(&self.terminal_error) { + Some(error) => Err(shared(&error)), + None => Ok(()), + } + } +} + +fn positive_bytes(bytes: u64, name: &str) -> Result { + usize::try_from(bytes) + .ok() + .filter(|bytes| *bytes > 0) + .with_context(|| InvalidWalObjectStoreSnafu { + reason: format!("{name} {bytes} is zero or too large"), + }) +} + +#[async_trait::async_trait] +impl LogStore for ObjectStoreLogStore { + type Error = Error; + + async fn stop(&self) -> Result<()> { + self.stopped.store(true, Ordering::Release); + let (response_tx, response_rx) = oneshot::channel(); + let sent = self + .command_tx + .send(Command::Stop { + response: response_tx, + }) + .await; + // A closed channel means the actor already exited. + if sent.is_ok() { + return response_rx.await.unwrap_or(Ok(())); + } + Ok(()) + } + + async fn append_batch(&self, _entries: Vec) -> Result { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn read( + &self, + _provider: &Provider, + _entry_id: EntryId, + _index: Option, + ) -> Result> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn create_namespace(&self, _ns: &Provider) -> Result<()> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn delete_namespace(&self, _ns: &Provider) -> Result<()> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn list_namespaces(&self) -> Result> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn obsolete( + &self, + _provider: &Provider, + _region_id: RegionId, + _entry_id: EntryId, + ) -> Result<()> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + fn entry( + &self, + _data: Vec, + _entry_id: EntryId, + _region_id: RegionId, + _provider: &Provider, + ) -> Result { + UnsupportedObjectStoreWalOperationSnafu.fail() + } + + fn latest_entry_id(&self, provider: &Provider) -> Result { + self.check_terminal()?; + let region_id = self.region_of(provider)?; + Ok(self.durable_entry_id_of(region_id)) + } +} + +enum Command { + Stop { + response: oneshot::Sender>, + }, +} + +struct Actor { + terminal_error: TerminalError, + stopped: Arc, + command_rx: mpsc::Receiver, + issued_entry_ids: HashMap, + next_object_seq: Option, + stop: Vec>>, +} + +impl Actor { + async fn run(mut self) { + while let Some(Command::Stop { response }) = self.command_rx.recv().await { + self.handle_stop(response); + if self.finish_stop() { + return; + } + } + } + + fn handle_stop(&mut self, response: oneshot::Sender>) { + self.stop.push(response); + } + + fn finish_stop(&mut self) -> bool { + if !self.is_stopped() || self.stop.is_empty() { + return false; + } + for response in self.stop.drain(..) { + let _ = response.send(Ok(())); + } + true + } + + fn is_stopped(&self) -> bool { + self.stopped.load(Ordering::Acquire) + } +} + +fn terminal(terminal_error: &TerminalError) -> Option> { + terminal_error + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() +} + +/// Records `error` as the terminal error unless one is already recorded, and +/// returns the recorded one. +fn set_terminal(terminal_error: &TerminalError, error: Error) -> Arc { + terminal_error + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get_or_insert_with(|| Arc::new(error)) + .clone() +} + +/// Wraps an error that several callers receive. +fn shared(error: &Arc) -> Error { + ObjectStoreWalSnafu.into_error(error.clone()) +} + +/// Rebuilds the catalog from object footers, +/// so recovery costs a few small reads per object however large the objects +/// are. Segments are not read; a segment checksum is verified by the read that +/// decodes it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects +/// at a time and indexed in sequence order, so the catalog checks the entry +/// ranges of every object against its predecessors like a sequential replay. +async fn recover(io: &dyn WalObjectIo) -> Result<(ObjectCatalog, u64, HashMap)> { + let objects = io.list().await?; + let mut catalog = ObjectCatalog::default(); + for (object, footer) in fetch_footers(io, objects, RECOVERY_CONCURRENCY).await? { + catalog + .insert_object(object.object_seq, footer) + .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + } + finish_recovery(catalog) +} + +fn finish_recovery( + catalog: ObjectCatalog, +) -> Result<(ObjectCatalog, u64, HashMap)> { + let next_object_seq = catalog.next_object_seq()?; + let durable_entry_ids = durable_entry_ids(&catalog); + Ok((catalog, next_object_seq, durable_entry_ids)) +} + +/// Fetches and verifies the footers of `objects`, up to `concurrency` objects +/// at a time, and returns them ordered by object sequence whatever the order +/// the fetches complete in. The first failure abandons the remaining fetches. +async fn fetch_footers( + io: &dyn WalObjectIo, + objects: Vec, + concurrency: usize, +) -> Result)>> { + let mut footers = futures::stream::iter(objects) + .map(|object| async move { + let footer = fetch_footer(io, &object).await?; + Ok((object, footer)) + }) + .buffer_unordered(concurrency) + .try_collect::>() + .await?; + footers.sort_unstable_by_key(|(object, _)| object.object_seq); + Ok(footers) +} + +/// Reads the header, trailer and footer of `object` and verifies them: the +/// header must carry the sequence of the key, the trailer must be well formed, +/// the footer must match the checksum the trailer holds and its segments must +/// tile the object body. +/// +/// A short object is read whole. Otherwise the header and a window at the end +/// of the object are read concurrently, and the footer is read separately only +/// when it starts before the window. +async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result> { + let ListedObject { + object_seq, size, .. + } = *object; + let invalid = |source: Error| { + InvalidWalObjectSnafu { + path: object.path.clone(), + } + .into_error(source) + }; + let object_len = usize::try_from(size) + .ok() + .filter(|len| *len >= MIN_OBJECT_LEN) + .with_context(|| CorruptedWalObjectSnafu { + reason: format!( + "truncated object, expected at least {MIN_OBJECT_LEN} bytes, actual {size}" + ), + }) + .map_err(invalid)?; + + let window = object_len.min(RECOVERY_TAIL_WINDOW); + let tail_start = object_len - window; + let (head, tail) = if tail_start == 0 { + let bytes = io.get(object_seq).await?; + (bytes.clone(), bytes) + } else { + futures::try_join!( + io.get_range(object_seq, 0, HEADER_LEN as u64), + io.get_range(object_seq, tail_start as u64, window as u64) + )? + }; + if head.len() < HEADER_LEN || tail.len() != window { + return Err(invalid( + CorruptedWalObjectSnafu { + reason: format!( + "object holds fewer bytes than the listed {size}, head {} bytes, tail {} bytes", + head.len(), + tail.len() + ), + } + .build(), + )); + } + + let (trailer, footer_range) = + locate_footer(object_seq, object_len, &head, &tail).map_err(invalid)?; + let footer = if footer_range.start >= tail_start { + tail.slice(footer_range.start - tail_start..footer_range.end - tail_start) + } else { + io.get_range( + object_seq, + footer_range.start as u64, + footer_range.len() as u64, + ) + .await? + }; + let footer = decode_footer(&footer, trailer).map_err(invalid)?; + verify_segment_ranges(&footer, footer_range.start).map_err(invalid)?; + Ok(footer) +} + +/// Verifies the header and trailer of the object `object_seq` of `object_len` +/// bytes from its first bytes `head` and its last bytes `tail`, and returns +/// the trailer with the range the footer occupies in the object. +fn locate_footer( + object_seq: u64, + object_len: usize, + head: &[u8], + tail: &[u8], +) -> Result<(FixedTrailer, Range)> { + let header = decode_header(head)?; + ensure!( + header.object_seq == object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "header sequence {} does not match key sequence {object_seq}", + header.object_seq + ), + } + ); + let trailer = decode_trailer(&tail[tail.len() - TRAILER_LEN..])?; + let footer_range = footer_range(trailer, object_len)?; + Ok((trailer, footer_range)) +} + +fn durable_entry_ids(catalog: &ObjectCatalog) -> HashMap { + let mut entry_ids = HashMap::new(); + for (_, footer) in catalog.objects_in_order() { + for entry in footer { + entry_ids + .entry(entry.region_id) + .and_modify(|current: &mut EntryId| *current = (*current).max(entry.max_entry_id)) + .or_insert(entry.max_entry_id); + } + } + entry_ids +} + +/// Object access of the store, so tests can inject failures. +#[async_trait::async_trait] +pub(crate) trait WalObjectIo: Send + Sync { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result; + + async fn get(&self, object_seq: u64) -> Result; + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result; + + async fn list(&self) -> Result>; + + fn object_path(&self, object_seq: u64) -> String; +} + +#[async_trait::async_trait] +impl WalObjectIo for ObjectStoreIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + ObjectStoreIo::put_if_absent(self, object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + ObjectStoreIo::get(self, object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + ObjectStoreIo::get_range(self, object_seq, offset, len).await + } + + async fn list(&self) -> Result> { + ObjectStoreIo::list(self).await + } + + fn object_path(&self, object_seq: u64) -> String { + ObjectStoreIo::object_path(self, object_seq) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use common_base::readable_size::ReadableSize; + use common_error::ext::{ErrorExt, RetryHint}; + use object_store::services::Memory; + use tokio::time::timeout; + + use super::*; + use crate::error::WalObjectStoreSnafu; + use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, entry_id}; + use crate::object_store_wal::format::{ + FOOTER_ENTRY_LEN, Header, Record, decode_object, encode_object, + }; + + const PREFIX: &str = "wal/datanodes/1/epochs/2"; + const WAIT: Duration = Duration::from_secs(30); + + fn memory_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap() + } + + fn config(flush_interval: Duration, max_batch_bytes: u64) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + storage_provider: String::new(), + flush_interval, + max_batch_bytes: ReadableSize(max_batch_bytes), + ..Default::default() + } + } + + fn eager() -> ObjectStoreWalConfig { + config(Duration::from_millis(10), 1) + } + + async fn open( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + ) -> Arc { + ObjectStoreLogStore::try_new(object_store, config, 1, 2) + .await + .unwrap() + } + + fn region(number: u32) -> RegionId { + RegionId::new(1, number) + } + + fn provider(region_id: RegionId) -> Provider { + Provider::object_store_provider(region_id, PREFIX.to_string()) + } + + fn latest(store: &ObjectStoreLogStore, region_id: RegionId) -> EntryId { + store.latest_entry_id(&provider(region_id)).unwrap() + } + + /// The id of the entry at `position` of its region in object `object_seq`. + fn id(object_seq: u64, position: u64) -> EntryId { + entry_id(object_seq, position) + } + #[tokio::test] + async fn test_store_rejects_invalid_config() { + for config in [ + config(Duration::from_millis(9), 1), + config(Duration::from_secs(1), 0), + ObjectStoreWalConfig { + prefix: "/absolute".to_string(), + ..config(Duration::from_secs(1), 1) + }, + ] { + let error = ObjectStoreLogStore::try_new(memory_store(), &config, 1, 2) + .await + .err() + .unwrap(); + assert!( + matches!(error, Error::InvalidWalObjectStore { .. }), + "unexpected error for {config:?}: {error:?}" + ); + } + } + + #[tokio::test] + async fn test_store_recovers_only_its_node_and_generation() { + let object_store = memory_store(); + let config = ObjectStoreWalConfig::default(); + let identities = [(1, 2), (3, 2), (1, 4)]; + for (index, (node_id, generation)) in identities.iter().enumerate() { + let encoded = encode_object( + Header { + object_seq: 0, + writer_instance: [0; 16], + }, + &[Record { + region_id: region(1), + entry_id: index as u64 + 1, + payload: Bytes::from_static(b"entry"), + }], + ) + .unwrap(); + ObjectStoreIo::new( + object_store.clone(), + config.node_prefix(*node_id, *generation), + ) + .unwrap() + .put_if_absent(0, encoded.bytes) + .await + .unwrap(); + } + // A root-level object must not be part of any node's recovery. + let root_io = ObjectStoreIo::new(object_store.clone(), &config.prefix).unwrap(); + root_io + .put_if_absent(0, Bytes::from_static(b"invalid")) + .await + .unwrap(); + for (index, (node_id, generation)) in identities.iter().enumerate() { + let store = + ObjectStoreLogStore::try_new(object_store.clone(), &config, *node_id, *generation) + .await + .unwrap(); + for (other_index, (other_node, other_generation)) in identities.iter().enumerate() { + let provider = Provider::object_store_provider( + region(1), + config.node_prefix(*other_node, *other_generation), + ); + if index == other_index { + assert_eq!(store.latest_entry_id(&provider).unwrap(), index as u64 + 1); + } else { + assert!(matches!( + store.latest_entry_id(&provider), + Err(Error::MismatchedWalPrefix { .. }) + )); + } + } + let root_provider = Provider::object_store_provider(region(1), config.prefix.clone()); + assert!(matches!( + store.latest_entry_id(&root_provider), + Err(Error::MismatchedWalPrefix { .. }) + )); + store.stop().await.unwrap(); + } + } + + /// Rebuilds the catalog by decoding whole objects as a recovery oracle. + async fn recover_by_decoding( + io: &dyn WalObjectIo, + ) -> Result<(ObjectCatalog, u64, HashMap)> { + let mut catalog = ObjectCatalog::default(); + for ListedObject { + object_seq, path, .. + } in io.list().await? + { + let bytes = io.get(object_seq).await?; + decode_object(&bytes) + .and_then(|decoded| { + ensure!( + decoded.header.object_seq == object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "header sequence {} does not match key sequence {object_seq}", + decoded.header.object_seq + ), + } + ); + catalog.insert_object(object_seq, decoded.footer) + }) + .with_context(|_| InvalidWalObjectSnafu { path })?; + } + finish_recovery(catalog) + } + + fn catalog_contents(catalog: &ObjectCatalog) -> Vec<(u64, Vec)> { + catalog + .objects_in_order() + .map(|(object_seq, footer)| (object_seq, footer.to_vec())) + .collect() + } + + async fn populate(object_store: &ObjectStore, objects: usize, regions: u32) { + for object in 0..objects { + let records = (1..=regions) + .filter(|number| !(object + *number as usize).is_multiple_of(3)) + .map(|number| Record { + region_id: region(number), + entry_id: id(object as u64, 1), + payload: Bytes::from(format!("o{object}-r{number}")), + }) + .collect::>(); + put_records(object_store, object as u64, &records).await; + } + } + + async fn put_records(object_store: &ObjectStore, object_seq: u64, records: &[Record]) { + let encoded = encode_object( + Header { + object_seq, + writer_instance: [0; 16], + }, + records, + ) + .unwrap(); + ObjectStoreIo::new(object_store.clone(), PREFIX) + .unwrap() + .put_if_absent(object_seq, encoded.bytes) + .await + .unwrap(); + } + + fn object_path(object_store: &ObjectStore, object_seq: u64) -> String { + ObjectStoreIo::new(object_store.clone(), PREFIX) + .unwrap() + .object_path(object_seq) + } + + async fn corrupt_object( + object_store: &ObjectStore, + path: &str, + corrupt: impl FnOnce(&mut Vec), + ) { + let mut bytes = object_store.read(path).await.unwrap().to_vec(); + corrupt(&mut bytes); + object_store.write(path, bytes).await.unwrap(); + } + + fn footer_of(bytes: &[u8]) -> (FixedTrailer, Vec) { + let trailer = decode_trailer(&bytes[bytes.len() - TRAILER_LEN..]).unwrap(); + let footer = + decode_footer(&bytes[footer_range(trailer, bytes.len()).unwrap()], trailer).unwrap(); + (trailer, footer) + } + + fn assert_invalid_object(error: &Error, path: &str, reason: &str) { + match error { + Error::InvalidWalObject { + path: actual, + source, + .. + } => { + assert_eq!(path, actual); + match &**source { + Error::CorruptedWalObject { reason: actual, .. } => assert!( + actual.contains(reason), + "expected reason to contain {reason:?}, actual {actual:?}" + ), + other => panic!("expected a corrupted object error, actual {other:?}"), + } + } + other => panic!("expected an invalid object error, actual {other:?}"), + } + } + + #[tokio::test] + async fn test_store_footer_recovery_matches_full_decode() { + let object_store = memory_store(); + populate(&object_store, 40, 5).await; + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + + let (catalog, next_object_seq, durable) = recover(&io).await.unwrap(); + let (expected_catalog, expected_next_object_seq, expected_durable) = + recover_by_decoding(&io).await.unwrap(); + + assert_eq!(40, catalog_contents(&catalog).len()); + assert_eq!( + catalog_contents(&expected_catalog), + catalog_contents(&catalog) + ); + assert_eq!(expected_next_object_seq, next_object_seq); + assert_eq!(40, next_object_seq); + assert_eq!(expected_durable, durable); + assert_eq!(5, durable.len()); + + let store = open(object_store, &eager()).await; + for number in 1..=5 { + let region_id = region(number); + // The objects `populate` gave the region one entry each. + let objects = (0..40u64) + .filter(|object| !(object + number as u64).is_multiple_of(3)) + .collect::>(); + assert_eq!(id(*objects.last().unwrap(), 1), durable[®ion_id]); + assert_eq!(durable[®ion_id], latest(&store, region_id)); + } + } + + #[tokio::test] + async fn test_store_recovery_rejects_corrupted_trailer_version_and_footer() { + type Corrupt = fn(&mut Vec); + let cases: &[(&str, Corrupt)] = &[ + ("truncated object", |bytes| { + bytes.truncate(MIN_OBJECT_LEN - 1) + }), + ("invalid header magic", |bytes| bytes[0] ^= 1), + ("invalid footer range", |bytes| { + let start = bytes.len() - TRAILER_LEN; + bytes[start..start + 8].copy_from_slice(&0u64.to_be_bytes()); + }), + ("overflows the object", |bytes| { + let start = bytes.len() - TRAILER_LEN; + bytes[start..start + 8].copy_from_slice(&u64::MAX.to_be_bytes()); + }), + ("invalid trailer magic", |bytes| { + let last = bytes.len() - 1; + bytes[last] ^= 1; + }), + ("unsupported format version 2", |bytes| { + bytes[8..10].copy_from_slice(&2u16.to_be_bytes()); + }), + ("footer checksum mismatch", |bytes| { + let (trailer, _) = footer_of(bytes); + bytes[trailer.footer_offset as usize] ^= 1; + }), + ]; + for (reason, corrupt) in cases { + let object_store = memory_store(); + populate(&object_store, 3, 2).await; + let path = object_path(&object_store, 1); + corrupt_object(&object_store, &path, *corrupt).await; + + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert_invalid_object(&error, &path, reason); + } + } + + #[tokio::test] + async fn test_store_recovery_rejects_header_sequence_mismatch() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let encoded = encode_object( + Header { + object_seq: 0, + writer_instance: [0; 16], + }, + &[Record { + region_id: region(1), + entry_id: 1, + payload: Bytes::from_static(b"a1"), + }], + ) + .unwrap(); + io.put_if_absent(5, encoded.bytes).await.unwrap(); + + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert_invalid_object( + &error, + &io.object_path(5), + "header sequence 0 does not match key sequence 5", + ); + } + + /// Overwrites the byte range of footer entry `index` and refreshes the + /// footer checksum, so the footer is intact but describes the wrong bytes. + fn rewrite_segment_range(bytes: &mut [u8], index: usize, offset: u64, len: u64) { + let trailer_start = bytes.len() - TRAILER_LEN; + let (trailer, _) = footer_of(bytes); + let footer = footer_range(trailer, bytes.len()).unwrap(); + let entry = footer.start + 4 + index * FOOTER_ENTRY_LEN; + bytes[entry + 28..entry + 36].copy_from_slice(&offset.to_be_bytes()); + bytes[entry + 36..entry + 44].copy_from_slice(&len.to_be_bytes()); + let checksum = crc32fast::hash(&bytes[footer]); + bytes[trailer_start + 16..trailer_start + 20].copy_from_slice(&checksum.to_be_bytes()); + } + + #[tokio::test] + async fn test_store_recovery_rejects_segment_ranges_that_do_not_tile_the_object() { + type Corrupt = fn(&mut Vec, &[FooterEntry]); + let cases: [(&str, Corrupt); 5] = [ + ("overflows the object", |bytes, footer| { + rewrite_segment_range(bytes, 0, u64::MAX, footer[0].segment_len); + }), + ("invalid segment range", |bytes, footer| { + let second = &footer[1]; + rewrite_segment_range(bytes, 1, second.segment_offset + 1, second.segment_len); + }), + ("invalid segment range", |bytes, footer| { + let second = &footer[1]; + rewrite_segment_range(bytes, 1, second.segment_offset - 1, second.segment_len); + }), + ("invalid segment range", |bytes, footer| { + let second = &footer[1]; + rewrite_segment_range(bytes, 1, second.segment_offset, second.segment_len + 1); + }), + ("segments end at", |bytes, footer| { + let second = &footer[1]; + rewrite_segment_range(bytes, 1, second.segment_offset, second.segment_len - 1); + }), + ]; + for (reason, corrupt) in cases { + let object_store = memory_store(); + populate(&object_store, 2, 3).await; + let path = object_path(&object_store, 1); + corrupt_object(&object_store, &path, |bytes| { + let (_, footer) = footer_of(bytes); + assert_eq!(2, footer.len()); + corrupt(bytes, &footer); + // The footer itself still verifies. + footer_of(bytes); + }) + .await; + + let error = ObjectStoreLogStore::try_new(object_store.clone(), &eager(), 1, 2) + .await + .unwrap_err(); + assert_invalid_object(&error, &path, reason); + let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); + let error = recover_by_decoding(&io).await.unwrap_err(); + assert_invalid_object(&error, &path, reason); + } + } + #[tokio::test] + async fn test_store_recovery_fetches_a_footer_longer_than_the_tail_window() { + let object_store = memory_store(); + let regions = (RECOVERY_TAIL_WINDOW / FOOTER_ENTRY_LEN + 100) as u32; + let records = (1..=regions) + .map(|number| Record { + region_id: region(number), + entry_id: 1, + payload: Bytes::from_static(b"wide"), + }) + .collect::>(); + put_records(&object_store, 0, &records).await; + put_object(&object_store, 1, region(1), &[id(1, 1)]).await; + + let (io, reads) = RecordingIo::over(object_store.clone()); + let objects = io.list().await.unwrap(); + let wide = &objects[0]; + let (trailer, _) = footer_of(&object_store.read(&wide.path).await.unwrap().to_vec()); + assert!(trailer.footer_len > RECOVERY_TAIL_WINDOW as u64); + + let (catalog, next_object_seq, durable) = recover(io.as_ref()).await.unwrap(); + let (expected_catalog, expected_next_object_seq, expected_durable) = + recover_by_decoding(io.as_ref()).await.unwrap(); + assert_eq!( + catalog_contents(&expected_catalog), + catalog_contents(&catalog) + ); + assert_eq!(expected_next_object_seq, next_object_seq); + assert_eq!(expected_durable, durable); + assert_eq!(regions as usize, durable.len()); + + // The wide object took the header, the tail window and the footer; + // the narrow one was read whole. + let mut wide_reads = reads + .lock() + .unwrap() + .iter() + .filter(|(object_seq, _, _)| *object_seq == wide.object_seq) + .map(|(_, offset, len)| (*offset, *len)) + .collect::>(); + wide_reads.sort_unstable(); + assert_eq!( + vec![ + (0, HEADER_LEN as u64), + (trailer.footer_offset, trailer.footer_len), + ( + wide.size - RECOVERY_TAIL_WINDOW as u64, + RECOVERY_TAIL_WINDOW as u64 + ), + ], + wide_reads + ); + assert!( + reads + .lock() + .unwrap() + .iter() + .all(|(object_seq, _, _)| *object_seq == wide.object_seq) + ); + + let store = open(object_store, &eager()).await; + assert_eq!(id(1, 1), latest(&store, region(1))); + assert_eq!(1, latest(&store, region(regions))); + store.stop().await.unwrap(); + } + + async fn put_object( + object_store: &ObjectStore, + object_seq: u64, + region_id: RegionId, + entry_ids: &[EntryId], + ) { + let records = entry_ids + .iter() + .map(|entry_id| Record { + region_id, + entry_id: *entry_id, + payload: Bytes::from(format!("e{entry_id}")), + }) + .collect::>(); + let encoded = encode_object( + Header { + object_seq, + writer_instance: [0; 16], + }, + &records, + ) + .unwrap(); + ObjectStoreIo::new(object_store.clone(), PREFIX) + .unwrap() + .put_if_absent(object_seq, encoded.bytes) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_store_resumes_sequence_and_durable_ids() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let (_, next, durable) = recover(&io).await.unwrap(); + assert_eq!(0, next); + assert!(durable.is_empty()); + + put_object(&object_store, 2, region(1), &[1, id(5, 7)]).await; + put_object(&object_store, 4, region(2), &[8]).await; + let (_, next, durable) = recover(&io).await.unwrap(); + assert_eq!(6, next); + assert_eq!( + HashMap::from([(region(1), id(5, 7)), (region(2), 8)]), + durable + ); + let store = open(object_store, &eager()).await; + assert_eq!(id(5, 7), latest(&store, region(1))); + assert_eq!(8, latest(&store, region(2))); + assert_eq!(0, latest(&store, region(3))); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_rejects_exhausted_sequence() { + for (object_seq, entry_id) in [(u64::MAX, 1), (OBJECT_SEQ_LIMIT - 1, 1), (0, u64::MAX)] { + let object_store = memory_store(); + put_object(&object_store, object_seq, region(1), &[entry_id]).await; + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalObjectSequenceExhausted { .. }), + "{error:?}" + ); + } + } + + #[tokio::test] + async fn test_store_recovery_rejects_conflicting_entry_ranges() { + for entries in [[1, 2], [2, 3]] { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1, 2]).await; + put_object(&object_store, 1, region(1), &entries).await; + let error = ObjectStoreLogStore::try_new(object_store.clone(), &eager(), 1, 2) + .await + .unwrap_err(); + assert_invalid_object(&error, &object_path(&object_store, 1), "entry"); + } + } + + #[tokio::test] + async fn test_store_rejects_foreign_providers() { + let store = open(memory_store(), &eager()).await; + let other = Provider::object_store_provider(region(1), "other/prefix".to_string()); + assert!(matches!(store.latest_entry_id(&other).unwrap_err(), + Error::MismatchedWalPrefix { expected, actual, .. } + if expected == PREFIX && actual == "other/prefix")); + let raft = Provider::raft_engine_provider(region(1).as_u64()); + assert!(matches!( + store.latest_entry_id(&raft), + Err(Error::InvalidProvider { .. }) + )); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_stop_is_idempotent() { + let store = open(memory_store(), &eager()).await; + let (first, second) = tokio::join!(store.stop(), store.stop()); + first.unwrap(); + second.unwrap(); + timeout(WAIT, store.command_tx.closed()).await.unwrap(); + store.stop().await.unwrap(); + assert!(store.stopped.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn test_store_stop_with_actor_gone() { + let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + drop(command_rx); + let store = ObjectStoreLogStore { + prefix: PREFIX.to_string(), + io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), + catalog: Arc::default(), + terminal_error: Arc::default(), + stopped: Arc::new(AtomicBool::new(false)), + command_tx, + }; + store.stop().await.unwrap(); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_actor_exits_when_the_store_is_dropped() { + let store = open(memory_store(), &eager()).await; + let io = Arc::downgrade(&store.io); + let stopped = Arc::downgrade(&store.stopped); + drop(store); + timeout(WAIT, async { + while stopped.strong_count() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(io.upgrade().is_none()); + } + + #[tokio::test] + async fn test_store_retains_first_terminal_error() { + let store = open(memory_store(), &eager()).await; + let first = set_terminal( + &store.terminal_error, + CorruptedWalObjectSnafu { reason: "first" }.build(), + ); + let second = set_terminal( + &store.terminal_error, + CorruptedWalObjectSnafu { reason: "second" }.build(), + ); + assert!(Arc::ptr_eq(&first, &second)); + match store.latest_entry_id(&provider(region(1))).unwrap_err() { + Error::ObjectStoreWal { source, .. } => assert!(Arc::ptr_eq(&first, &source)), + error => panic!("{error:?}"), + } + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_bounds_concurrency_and_orders_by_sequence() { + let object_store = memory_store(); + populate(&object_store, 16, 3).await; + let (io, mut parked) = ParkedIo::over(object_store); + let objects = io.list().await.unwrap(); + + let mut fetch = { + let io = io.clone(); + tokio::spawn(async move { + fetch_footers(io.as_ref(), objects, RECOVERY_CONCURRENCY) + .await + .unwrap() + .into_iter() + .map(|(object, _)| object.object_seq) + .collect::>() + }) + }; + let mut wave = Vec::new(); + for _ in 0..RECOVERY_CONCURRENCY { + wave.push(timeout(WAIT, parked.recv()).await.unwrap().unwrap()); + } + assert_eq!( + (0..8).collect::>(), + wave.iter() + .map(|(object_seq, _)| *object_seq) + .collect::>() + ); + // No ninth fetch can start while eight are parked. + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(parked.try_recv().is_err()); + + // Complete the initial wave in reverse order, parking each replacement. + let mut replacements = Vec::new(); + for (expected_next, (_, release)) in (8..16).zip(wave.into_iter().rev()) { + release.send(true).unwrap(); + let (object_seq, release) = timeout(WAIT, parked.recv()).await.unwrap().unwrap(); + assert_eq!(expected_next, object_seq); + assert!(!fetch.is_finished()); + replacements.push(release); + } + for release in replacements { + release.send(true).unwrap(); + } + assert_eq!( + (0..16).collect::>(), + timeout(WAIT, &mut fetch).await.unwrap().unwrap() + ); + assert_eq!(8, io.max_in_flight.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_store_recovery_abandons_pending_fetches_on_failure() { + let object_store = memory_store(); + populate(&object_store, 16, 3).await; + let (io, mut parked) = ParkedIo::over(object_store.clone()); + let fetch = { + let io = io.clone(); + tokio::spawn(async move { + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + }) + }; + let mut wave = Vec::new(); + for _ in 0..RECOVERY_CONCURRENCY { + wave.push(timeout(WAIT, parked.recv()).await.unwrap().unwrap()); + } + let (failed_seq, release) = wave.pop().unwrap(); + release.send(false).unwrap(); + let error = timeout(WAIT, fetch).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "read", path, .. } + if path == &io.object_path(failed_seq)) + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + for (_, release) in wave { + assert!(release.is_closed()); + } + assert!(parked.try_recv().is_err()); + let store = open(object_store, &eager()).await; + assert_eq!(id(15, 1), latest(&store, region(1))); + assert_eq!(16, store.catalog.read().unwrap().next_object_seq().unwrap()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_reads_only_header_and_tail_of_large_object() { + let object_store = memory_store(); + put_records( + &object_store, + 0, + &[Record { + region_id: region(1), + entry_id: 1, + payload: Bytes::from(vec![0; RECOVERY_TAIL_WINDOW * 2]), + }], + ) + .await; + let path = object_path(&object_store, 0); + corrupt_object(&object_store, &path, |bytes| bytes[HEADER_LEN + 20] ^= 1).await; + let (io, reads) = RecordingIo::over(object_store); + let object = io.list().await.unwrap().remove(0); + let store = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap(); + assert_eq!(1, latest(&store, region(1))); + let mut reads = reads.lock().unwrap().clone(); + reads.sort_unstable(); + assert_eq!( + vec![ + (0, 0, HEADER_LEN as u64), + ( + 0, + object.size - RECOVERY_TAIL_WINDOW as u64, + RECOVERY_TAIL_WINDOW as u64 + ) + ], + reads + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_rejects_short_read() { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1]).await; + let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); + let mut object = io.list().await.unwrap().remove(0); + object.size += 1; + let error = fetch_footer(&io, &object).await.unwrap_err(); + assert_invalid_object(&error, &object.path, "fewer bytes than the listed"); + } + + #[tokio::test] + async fn test_store_unsupported_operations() { + let store = open(memory_store(), &eager()).await; + let region = region(1); + let provider = provider(region); + let errors = [ + store.append_batch(Vec::new()).await.unwrap_err(), + store.read(&provider, 0, None).await.err().unwrap(), + store.create_namespace(&provider).await.unwrap_err(), + store.delete_namespace(&provider).await.unwrap_err(), + store.list_namespaces().await.unwrap_err(), + store.obsolete(&provider, region, 1).await.unwrap_err(), + store.obsolete_all(&provider, region).await.unwrap_err(), + store.entry(Vec::new(), 1, region, &provider).unwrap_err(), + ]; + for error in errors { + assert!(matches!( + error, + Error::UnsupportedObjectStoreWalOperation { .. } + )); + assert_eq!( + common_error::status_code::StatusCode::Unsupported, + error.status_code() + ); + } + store.stop().await.unwrap(); + } + + struct ParkedIo { + inner: ObjectStoreIo, + parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, + in_flight: AtomicUsize, + max_in_flight: AtomicUsize, + } + + impl ParkedIo { + fn over( + object_store: ObjectStore, + ) -> ( + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + let (parked, parked_rx) = mpsc::unbounded_channel(); + ( + Arc::new(Self { + inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), + parked, + in_flight: AtomicUsize::new(0), + max_in_flight: AtomicUsize::new(0), + }), + parked_rx, + ) + } + } + + #[async_trait::async_trait] + impl WalObjectIo for ParkedIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + self.inner.put_if_absent(object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); + let (release, released) = oneshot::channel(); + self.parked.send((object_seq, release)).unwrap(); + let success = released.await.unwrap(); + self.in_flight.fetch_sub(1, Ordering::SeqCst); + if !success { + return Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected failure", + ) + .set_temporary()) + .context(WalObjectStoreSnafu { + operation: "read", + path: self.object_path(object_seq), + }); + } + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.inner.get_range(object_seq, offset, len).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } + + type RangeReads = Arc>>; + + /// Object access that records every range read as (sequence, offset, length). + struct RecordingIo { + inner: ObjectStoreIo, + reads: RangeReads, + } + + impl RecordingIo { + fn over(object_store: ObjectStore) -> (Arc, RangeReads) { + let reads = RangeReads::default(); + let io = Self { + inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), + reads: reads.clone(), + }; + (Arc::new(io), reads) + } + } + + #[async_trait::async_trait] + impl WalObjectIo for RecordingIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + self.inner.put_if_absent(object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.reads.lock().unwrap().push((object_seq, offset, len)); + self.inner.get_range(object_seq, offset, len).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } +} From fbbc017be44072132e779c4a9e3a433efbe7032a Mon Sep 17 00:00:00 2001 From: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> Date: Sat, 19 Sep 2026 02:26:31 +0000 Subject: [PATCH 11/63] feat: expose region min/max timestamp in region_statistics (#9060) * feat: expose region min/max timestamp in region_statistics Signed-off-by: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> * test: cover region_statistic time range assembly and projected values Signed-off-by: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> --------- Signed-off-by: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> --- .../information_schema/region_statistics.rs | 39 +++++++- src/common/meta/src/datanode.rs | 12 ++- src/meta-srv/src/gc/handler.rs | 2 + src/meta-srv/src/gc/mock.rs | 2 + .../handler/collect_leader_region_handler.rs | 2 + src/meta-srv/src/handler/failure_handler.rs | 2 + .../src/handler/persist_stats_handler.rs | 2 + .../src/handler/region_lease_handler.rs | 2 + src/meta-srv/src/selector/weight_compute.rs | 6 ++ src/metric-engine/src/utils.rs | 76 ++++++++++++++++ src/mito2/src/engine/basic_test.rs | 88 +++++++++++++++++++ src/mito2/src/memtable/version.rs | 12 +++ src/mito2/src/region.rs | 9 ++ src/mito2/src/sst/version.rs | 80 ++++++++++++++++- src/standalone/src/information_extension.rs | 2 + src/store-api/src/region_engine.rs | 8 ++ .../region_statistics.result | 57 ++++++++++++ .../information_schema/region_statistics.sql | 35 ++++++++ .../common/system/information_schema.result | 2 + 19 files changed, 435 insertions(+), 3 deletions(-) diff --git a/src/catalog/src/system_schema/information_schema/region_statistics.rs b/src/catalog/src/system_schema/information_schema/region_statistics.rs index b781c277532..8d1a15ddcc3 100644 --- a/src/catalog/src/system_schema/information_schema/region_statistics.rs +++ b/src/catalog/src/system_schema/information_schema/region_statistics.rs @@ -20,13 +20,18 @@ use common_error::ext::BoxedError; use common_meta::datanode::RegionStat; use common_recordbatch::adapter::RecordBatchStreamAdapter; use common_recordbatch::{DfSendableRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_time::timestamp::TimeUnit; use datafusion::execution::TaskContext; use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter; use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream; use datatypes::prelude::{ConcreteDataType, ScalarVectorBuilder, VectorRef}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::timestamp::TimestampMillisecond; use datatypes::value::Value; -use datatypes::vectors::{StringVectorBuilder, UInt32VectorBuilder, UInt64VectorBuilder}; +use datatypes::vectors::{ + StringVectorBuilder, TimestampMillisecondVectorBuilder, UInt32VectorBuilder, + UInt64VectorBuilder, +}; use snafu::ResultExt; use store_api::storage::{ScanRequest, TableId}; @@ -51,6 +56,8 @@ const SST_NUM: &str = "sst_num"; const INDEX_SIZE: &str = "index_size"; const ENGINE: &str = "engine"; const REGION_ROLE: &str = "region_role"; +const MIN_TIMESTAMP: &str = "min_timestamp"; +const MAX_TIMESTAMP: &str = "max_timestamp"; const INIT_CAPACITY: usize = 42; @@ -109,6 +116,16 @@ impl InformationSchemaRegionStatistics { ColumnSchema::new(INDEX_SIZE, ConcreteDataType::uint64_datatype(), true), ColumnSchema::new(ENGINE, ConcreteDataType::string_datatype(), true), ColumnSchema::new(REGION_ROLE, ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + MIN_TIMESTAMP, + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ), + ColumnSchema::new( + MAX_TIMESTAMP, + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ), ])) } @@ -175,6 +192,8 @@ struct InformationSchemaRegionStatisticsBuilder { index_sizes: UInt64VectorBuilder, engines: StringVectorBuilder, region_roles: StringVectorBuilder, + min_timestamps: TimestampMillisecondVectorBuilder, + max_timestamps: TimestampMillisecondVectorBuilder, } impl InformationSchemaRegionStatisticsBuilder { @@ -197,6 +216,8 @@ impl InformationSchemaRegionStatisticsBuilder { index_sizes: UInt64VectorBuilder::with_capacity(INIT_CAPACITY), engines: StringVectorBuilder::with_capacity(INIT_CAPACITY), region_roles: StringVectorBuilder::with_capacity(INIT_CAPACITY), + min_timestamps: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY), + max_timestamps: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY), } } @@ -261,6 +282,20 @@ impl InformationSchemaRegionStatisticsBuilder { self.index_sizes.push(Some(region_stat.index_size)); self.engines.push(Some(®ion_stat.engine)); self.region_roles.push(Some(®ion_stat.role.to_string())); + // Floor the min and ceil the max so the window stays a superset: a narrower + // one would hide regions from a time-bounded lookup. + self.min_timestamps.push( + region_stat + .min_timestamp + .and_then(|ts| ts.convert_to(TimeUnit::Millisecond)) + .map(TimestampMillisecond), + ); + self.max_timestamps.push( + region_stat + .max_timestamp + .and_then(|ts| ts.convert_to_ceil(TimeUnit::Millisecond)) + .map(TimestampMillisecond), + ); } fn finish(&mut self) -> Result { @@ -280,6 +315,8 @@ impl InformationSchemaRegionStatisticsBuilder { Arc::new(self.index_sizes.finish()), Arc::new(self.engines.finish()), Arc::new(self.region_roles.finish()), + Arc::new(self.min_timestamps.finish()), + Arc::new(self.max_timestamps.finish()), ]; RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu) diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index d50c37c381b..81135aefe0e 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -16,7 +16,7 @@ use std::collections::{HashMap, HashSet}; use std::str::FromStr; use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, RequestHeader}; -use common_time::util as time_util; +use common_time::{Timestamp, util as time_util}; use lazy_static::lazy_static; use regex::Regex; use serde::{Deserialize, Serialize}; @@ -117,6 +117,12 @@ pub struct RegionStat { /// **Only used by remote WAL prune.** /// In mito engine, this is the same as `data_topic_latest_entry_id`. pub metadata_topic_latest_entry_id: u64, + /// The earliest timestamp of the region's current data, if it holds any. + #[serde(default)] + pub min_timestamp: Option, + /// The latest timestamp of the region's current data, if it holds any. + #[serde(default)] + pub max_timestamp: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -331,6 +337,8 @@ impl From<&api::v1::meta::RegionStat> for RegionStat { query_scanned_bytes: region_stat.query_scanned_bytes, data_topic_latest_entry_id: region_stat.data_topic_latest_entry_id, metadata_topic_latest_entry_id: region_stat.metadata_topic_latest_entry_id, + min_timestamp: region_stat.min_timestamp, + max_timestamp: region_stat.max_timestamp, } } } @@ -593,6 +601,8 @@ mod tests { query_scanned_bytes: 20, data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, }], ..Default::default() }; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index eda87fbb144..960a05eebce 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -499,5 +499,7 @@ fn dropped_region_stat(region_id: RegionId) -> RegionStat { query_scanned_bytes: 0, data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, } } diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index dcf913462c2..e0409de6e40 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -450,6 +450,8 @@ fn mock_region_stat( index_size: 0, data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, diff --git a/src/meta-srv/src/handler/collect_leader_region_handler.rs b/src/meta-srv/src/handler/collect_leader_region_handler.rs index 2b6aa1cb9f1..6ae753e4808 100644 --- a/src/meta-srv/src/handler/collect_leader_region_handler.rs +++ b/src/meta-srv/src/handler/collect_leader_region_handler.rs @@ -88,6 +88,8 @@ mod tests { index_size: 0, data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, diff --git a/src/meta-srv/src/handler/failure_handler.rs b/src/meta-srv/src/handler/failure_handler.rs index fa089c14c2a..5f1448d3ab6 100644 --- a/src/meta-srv/src/handler/failure_handler.rs +++ b/src/meta-srv/src/handler/failure_handler.rs @@ -109,6 +109,8 @@ mod tests { written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, } } acc.stat = Some(Stat { diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index fcdb9fa04bc..04d63e26e2f 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -302,6 +302,8 @@ mod tests { query_scanned_bytes: 0, data_topic_latest_entry_id: 200, metadata_topic_latest_entry_id: 200, + min_timestamp: None, + max_timestamp: None, } } diff --git a/src/meta-srv/src/handler/region_lease_handler.rs b/src/meta-srv/src/handler/region_lease_handler.rs index c8d3e75264f..98ddd39cc86 100644 --- a/src/meta-srv/src/handler/region_lease_handler.rs +++ b/src/meta-srv/src/handler/region_lease_handler.rs @@ -182,6 +182,8 @@ mod test { written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, } } diff --git a/src/meta-srv/src/selector/weight_compute.rs b/src/meta-srv/src/selector/weight_compute.rs index dfbabfa650b..27dcdf99b97 100644 --- a/src/meta-srv/src/selector/weight_compute.rs +++ b/src/meta-srv/src/selector/weight_compute.rs @@ -202,6 +202,8 @@ mod tests { written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, }], ..Default::default() } @@ -234,6 +236,8 @@ mod tests { written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, }], ..Default::default() } @@ -266,6 +270,8 @@ mod tests { written_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, }], ..Default::default() } diff --git a/src/metric-engine/src/utils.rs b/src/metric-engine/src/utils.rs index 1620b1bb9ef..503b75a1cc5 100644 --- a/src/metric-engine/src/utils.rs +++ b/src/metric-engine/src/utils.rs @@ -84,6 +84,10 @@ fn merge_region_statistic( query_scanned_bytes: data_stat.query_scanned_bytes, data_topic_latest_entry_id: data_stat.data_topic_latest_entry_id, metadata_topic_latest_entry_id: metadata_stat.metadata_topic_latest_entry_id, + // Metadata rows are written at timestamp 0, so merging their bounds would + // pin every metric region's minimum to the epoch. + min_timestamp: data_stat.min_timestamp, + max_timestamp: data_stat.max_timestamp, } } @@ -120,6 +124,8 @@ pub(crate) fn encode_manifest_info_to_extensions( #[cfg(test)] mod tests { + use common_time::Timestamp; + use super::*; #[test] @@ -144,6 +150,76 @@ mod tests { assert_eq!(to_data_region_id(region_id), expected_region_id); } + #[test] + fn merge_region_statistic_ignores_metadata_timestamps() { + // Metadata rows are written at timestamp 0; folding them in would report + // 1970 as the minimum for a region holding only modern samples. + let metadata_stat = RegionStatistic { + min_timestamp: Some(Timestamp::new_millisecond(0)), + max_timestamp: Some(Timestamp::new_millisecond(0)), + ..Default::default() + }; + let data_stat = RegionStatistic { + min_timestamp: Some(Timestamp::new_millisecond(1_700_000_000_000)), + max_timestamp: Some(Timestamp::new_millisecond(1_700_000_001_000)), + ..Default::default() + }; + + let statistic = merge_region_statistic(&metadata_stat, &data_stat); + + assert_eq!( + statistic.min_timestamp, + Some(Timestamp::new_millisecond(1_700_000_000_000)) + ); + assert_eq!( + statistic.max_timestamp, + Some(Timestamp::new_millisecond(1_700_000_001_000)) + ); + } + + #[test] + fn merge_region_statistic_reports_no_range_for_an_empty_data_region() { + // Metadata present but no samples: the region has no data to bound, so + // both ends must stay NULL rather than collapsing onto the epoch. + let metadata_stat = RegionStatistic { + min_timestamp: Some(Timestamp::new_millisecond(0)), + max_timestamp: Some(Timestamp::new_millisecond(0)), + ..Default::default() + }; + + let statistic = merge_region_statistic(&metadata_stat, &RegionStatistic::default()); + + assert_eq!(statistic.min_timestamp, None); + assert_eq!(statistic.max_timestamp, None); + } + + #[test] + fn merge_region_statistic_keeps_pre_epoch_samples() { + // Samples entirely before 1970 must not have their maximum pulled up to + // the metadata timestamp. + let metadata_stat = RegionStatistic { + min_timestamp: Some(Timestamp::new_millisecond(0)), + max_timestamp: Some(Timestamp::new_millisecond(0)), + ..Default::default() + }; + let data_stat = RegionStatistic { + min_timestamp: Some(Timestamp::new_millisecond(-2_000)), + max_timestamp: Some(Timestamp::new_millisecond(-1_000)), + ..Default::default() + }; + + let statistic = merge_region_statistic(&metadata_stat, &data_stat); + + assert_eq!( + statistic.min_timestamp, + Some(Timestamp::new_millisecond(-2_000)) + ); + assert_eq!( + statistic.max_timestamp, + Some(Timestamp::new_millisecond(-1_000)) + ); + } + #[test] fn merge_region_statistic_uses_data_region_query_stats() { let metadata_stat = RegionStatistic { diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index f4c843f1ce4..776e1f78abe 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -24,6 +24,7 @@ use common_base::readable_size::ReadableSize; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; +use common_time::Timestamp; use common_wal::options::WAL_OPTIONS_KEY; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; @@ -808,6 +809,93 @@ async fn test_region_usage_with_format(flat_format: bool) { assert!(region_stat.estimated_disk_size() > 3000); } +#[tokio::test] +async fn test_region_usage_time_range() { + test_region_usage_time_range_with_format(false).await; + test_region_usage_time_range_with_format(true).await; +} + +async fn test_region_usage_time_range_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("region_usage_time_range").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 column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + + // An empty region has nothing to bound. + let region_stat = region.region_statistic(); + assert_eq!(region_stat.min_timestamp, None); + assert_eq!(region_stat.max_timestamp, None); + + // Unflushed data lives only in the memtable, so the range must come from there. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("a", 10, 20, 0), + }, + ) + .await; + + let region_stat = region.region_statistic(); + assert_eq!( + region_stat.min_timestamp, + Some(Timestamp::new_millisecond(10_000)) + ); + assert_eq!( + region_stat.max_timestamp, + Some(Timestamp::new_millisecond(19_000)) + ); + + // After a flush the same range is served by the SST instead. + flush_region(&engine, region_id, None).await; + + let region_stat = region.region_statistic(); + assert_eq!( + region_stat.min_timestamp, + Some(Timestamp::new_millisecond(10_000)) + ); + assert_eq!( + region_stat.max_timestamp, + Some(Timestamp::new_millisecond(19_000)) + ); + + // A later write lands in the memtable below the SST range: the minimum must + // come from the memtable and the maximum from the SST, so a merge that drops + // either side is caught. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("b", 0, 5, 0), + }, + ) + .await; + + let region_stat = region.region_statistic(); + assert_eq!( + region_stat.min_timestamp, + Some(Timestamp::new_millisecond(0)) + ); + assert_eq!( + region_stat.max_timestamp, + Some(Timestamp::new_millisecond(19_000)) + ); +} + #[tokio::test] async fn test_engine_with_write_cache() { test_engine_with_write_cache_with_format(false).await; diff --git a/src/mito2/src/memtable/version.rs b/src/mito2/src/memtable/version.rs index 23ad623b896..d707128c93a 100644 --- a/src/mito2/src/memtable/version.rs +++ b/src/mito2/src/memtable/version.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use std::time::Duration; +use common_time::Timestamp; use smallvec::SmallVec; use store_api::metadata::RegionMetadataRef; use store_api::storage::SequenceNumber; @@ -162,6 +163,17 @@ impl MemtableVersion { + self.mutable.num_rows() } + /// Returns the time range covered by the memtables, if any hold data. + pub(crate) fn time_range(&self) -> Option<(Timestamp, Timestamp)> { + let mut mutables = Vec::new(); + self.mutable.list_memtables(&mut mutables); + self.immutables + .iter() + .chain(mutables.iter()) + .filter_map(|mem| mem.stats().time_range()) + .reduce(|(min_a, max_a), (min_b, max_b)| (min_a.min(min_b), max_a.max(max_b))) + } + /// Returns true if the memtable version is empty. /// /// The version is empty when mutable memtable is empty and there is no diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index dd392f5b0ab..521a0e08a21 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -701,6 +701,13 @@ impl MitoRegion { let manifest_version = self.stats.manifest_version(); let file_removed_cnt = self.stats.file_removed_cnt(); + let time_range = match (version.ssts.time_range(), version.memtables.time_range()) { + (Some((sst_min, sst_max)), Some((mem_min, mem_max))) => { + Some((sst_min.min(mem_min), sst_max.max(mem_max))) + } + (range, None) | (None, range) => range, + }; + let topic_latest_entry_id = self.topic_latest_entry_id.load(Ordering::Relaxed); let written_bytes = self.region_stats.written_bytes.load(Ordering::Relaxed); let query_cpu_time = self.region_stats.query_cpu_time.load(Ordering::Relaxed); @@ -727,6 +734,8 @@ impl MitoRegion { written_bytes, query_cpu_time, query_scanned_bytes, + min_timestamp: time_range.map(|(min, _)| min), + max_timestamp: time_range.map(|(_, max)| max), } } diff --git a/src/mito2/src/sst/version.rs b/src/mito2/src/sst/version.rs index 7c6e6fe2033..172d8701a73 100644 --- a/src/mito2/src/sst/version.rs +++ b/src/mito2/src/sst/version.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use common_time::{TimeToLive, Timestamp}; use store_api::storage::{FileId, RegionId}; -use crate::sst::file::{FileHandle, FileMeta, Level, MAX_LEVEL}; +use crate::sst::file::{FileHandle, FileMeta, FileTimeRange, Level, MAX_LEVEL}; use crate::sst::file_purger::FilePurgerRef; /// A version of all SSTs in a region. @@ -152,6 +152,16 @@ impl SstVersion { .sum() } + /// Returns the time range covered by every file in this version, including + /// files referenced from other regions after a repartition. + pub(crate) fn time_range(&self) -> Option { + self.levels + .iter() + .flat_map(|level_meta| level_meta.files.values()) + .map(|file_handle| file_handle.time_range()) + .reduce(|(min_a, max_a), (min_b, max_b)| (min_a.min(min_b), max_a.max(max_b))) + } + /// Returns the space occupied by SST data files owned by `region_id`. pub(crate) fn owned_sst_usage(&self, region_id: RegionId) -> u64 { self.levels @@ -257,6 +267,74 @@ mod tests { use super::*; use crate::test_util::new_noop_file_purger; + #[test] + fn time_range_spans_files_referenced_from_other_regions() { + let purger = new_noop_file_purger(); + let owned = FileMeta { + file_id: FileId::random(), + region_id: RegionId::new(1, 1), + time_range: ( + Timestamp::new_millisecond(200), + Timestamp::new_millisecond(300), + ), + ..Default::default() + }; + let referenced = FileMeta { + file_id: FileId::random(), + region_id: RegionId::new(2, 1), + time_range: ( + Timestamp::new_millisecond(50), + Timestamp::new_millisecond(100), + ), + ..Default::default() + }; + + let mut version = SstVersion::new(); + version.add_files(purger, [owned, referenced].into_iter()); + + assert_eq!( + version.time_range(), + Some(( + Timestamp::new_millisecond(50), + Timestamp::new_millisecond(300) + )) + ); + } + + #[test] + fn time_range_compares_across_units() { + let purger = new_noop_file_purger(); + let seconds = FileMeta { + file_id: FileId::random(), + time_range: (Timestamp::new_second(1), Timestamp::new_second(2)), + ..Default::default() + }; + let millis = FileMeta { + file_id: FileId::random(), + time_range: ( + Timestamp::new_millisecond(500), + Timestamp::new_millisecond(2500), + ), + ..Default::default() + }; + + let mut version = SstVersion::new(); + version.add_files(purger, [seconds, millis].into_iter()); + + assert_eq!( + version.time_range(), + Some(( + Timestamp::new_millisecond(500), + Timestamp::new_millisecond(2500) + )) + ); + } + + #[test] + fn time_range_is_none_without_files() { + assert_eq!(SstVersion::new().time_range(), None); + } + #[test] fn test_add_files() { let purger = new_noop_file_purger(); diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index c8be92a6f0a..8d52eb12ccc 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -139,6 +139,8 @@ impl InformationExtension for StandaloneInformationExtension { written_bytes: region_stat.written_bytes, query_cpu_time: region_stat.query_cpu_time, query_scanned_bytes: region_stat.query_scanned_bytes, + min_timestamp: region_stat.min_timestamp, + max_timestamp: region_stat.max_timestamp, } }) .collect::>(); diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 6a45d3db7d7..6577880b415 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -576,6 +576,14 @@ pub struct RegionStatistic { /// The total scanned bytes of the region since region opened. #[serde(default)] pub query_scanned_bytes: u64, + /// The earliest timestamp of the region's current data, or `None` if it holds + /// none. Unlike the size and row counters this covers files referenced from + /// other regions too, since a time range cannot double count. + #[serde(default)] + pub min_timestamp: Option, + /// The latest timestamp of the region's current data. See [`Self::min_timestamp`]. + #[serde(default)] + pub max_timestamp: Option, /// The latest entry id of the region's remote WAL since last flush. /// For metric engine, there're two latest entry ids, one for data and one for metadata. /// TODO(weny): remove this two fields and use single instead. diff --git a/tests/cases/standalone/common/information_schema/region_statistics.result b/tests/cases/standalone/common/information_schema/region_statistics.result index e2b89df6cf8..e6f9c9e47fc 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.result +++ b/tests/cases/standalone/common/information_schema/region_statistics.result @@ -49,3 +49,60 @@ DROP TABLE test; Affected Rows: 0 +-- The projected columns are millisecond-typed. A nanosecond time index proves +-- min_timestamp floors and max_timestamp ceils, so the reported window stays a +-- superset of the data and cannot hide a region from a time-bounded lookup. +CREATE TABLE precise ( + a int primary key, + ts timestamp(9) time index, +); + +Affected Rows: 0 + +-- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. +INSERT INTO precise VALUES + (1, 1500000), + (2, 2500000), + (3, 3999999); + +Affected Rows: 3 + +-- SQLNESS SLEEP 3s +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); + ++-------------------------+-------------------------+ +| min_timestamp | max_timestamp | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.001 | 1970-01-01T00:00:00.004 | ++-------------------------+-------------------------+ + +DROP TABLE precise; + +Affected Rows: 0 + +-- A region with no rows has nothing to bound, so both ends stay NULL rather +-- than collapsing onto the epoch. +CREATE TABLE empty_region ( + a int primary key, + ts timestamp time index, +); + +Affected Rows: 0 + +-- SQLNESS SLEEP 3s +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); + ++---------------+---------------+ +| min_timestamp | max_timestamp | ++---------------+---------------+ +| | | ++---------------+---------------+ + +DROP TABLE empty_region; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/information_schema/region_statistics.sql b/tests/cases/standalone/common/information_schema/region_statistics.sql index b5202d051c7..60264fbfb57 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.sql +++ b/tests/cases/standalone/common/information_schema/region_statistics.sql @@ -29,3 +29,38 @@ SELECT SUM(region_rows), SUM(written_bytes_since_open), SUM(query_cpu_time_milli SELECT data_length, index_length, avg_row_length, table_rows FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test'; DROP TABLE test; + +-- The projected columns are millisecond-typed. A nanosecond time index proves +-- min_timestamp floors and max_timestamp ceils, so the reported window stays a +-- superset of the data and cannot hide a region from a time-bounded lookup. +CREATE TABLE precise ( + a int primary key, + ts timestamp(9) time index, +); + +-- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. +INSERT INTO precise VALUES + (1, 1500000), + (2, 2500000), + (3, 3999999); + +-- SQLNESS SLEEP 3s +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); + +DROP TABLE precise; + +-- A region with no rows has nothing to bound, so both ends stay NULL rather +-- than collapsing onto the epoch. +CREATE TABLE empty_region ( + a int primary key, + ts timestamp time index, +); + +-- SQLNESS SLEEP 3s +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); + +DROP TABLE empty_region; diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index 7fdfed0defc..fa9dd02a782 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -380,7 +380,9 @@ order by table_schema, table_name, column_name; | greptime | information_schema | region_statistics | engine | 14 | 2147483647 | 2147483647 | | | | utf8 | utf8_bin | | | select,insert | | String | string | FIELD | | YES | string | | | | greptime | information_schema | region_statistics | index_size | 13 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | region_statistics | manifest_size | 10 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | +| greptime | information_schema | region_statistics | max_timestamp | 17 | | | | | 3 | | | | | select,insert | | TimestampMillisecond | timestamp(3) | FIELD | | YES | timestamp(3) | | | | greptime | information_schema | region_statistics | memtable_size | 9 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | +| greptime | information_schema | region_statistics | min_timestamp | 16 | | | | | 3 | | | | | select,insert | | TimestampMillisecond | timestamp(3) | FIELD | | YES | timestamp(3) | | | | greptime | information_schema | region_statistics | query_cpu_time_millis | 6 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | region_statistics | query_scanned_bytes | 7 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | region_statistics | region_id | 1 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | From 0e5bcaee549d3337f5da302a65bfbc7774c85039 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sat, 19 Sep 2026 02:46:12 +0000 Subject: [PATCH 12/63] ci: pin crate-ci/typos to v1.50.2 (#9246) crate-ci/typos renamed its default branch from master to main on 2026-09-18 (0a3d75e) and the master branch is gone, so every workflow run since then fails at job setup with: Unable to resolve action `crate-ci/typos@master`, unable to find version `master` Pin to the latest release tag instead of tracking a branch. This matches how the other actions in these two workflows are referenced and keeps the check reproducible. Signed-off-by: Dennis Zhuang --- .github/workflows/checks.yml | 2 +- .github/workflows/docs.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index e27c2d50c3e..6e591ba9dfd 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -39,7 +39,7 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: crate-ci/typos@master + - uses: crate-ci/typos@v1.50.2 - name: Check the config docs run: | make config-docs && \ diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 40ab1984278..3006dc7f954 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -45,7 +45,7 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: crate-ci/typos@master + - uses: crate-ci/typos@v1.50.2 license-header-check: if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} From 2ab8475b84bf2bbe2b864db56d753b534e2c50c4 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Sat, 19 Sep 2026 09:27:10 +0000 Subject: [PATCH 13/63] feat: add status label to datanode failed-insert metric (#9191) * feat: label datanode failed inserts by status Signed-off-by: dhruvxvaishnav * fix: rename failed-insert metric label to code Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav Co-authored-by: Weny Xu --- src/datanode/src/metrics.rs | 6 ++++-- src/datanode/src/region_server.rs | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/datanode/src/metrics.rs b/src/datanode/src/metrics.rs index 74da5b153f4..7573490c747 100644 --- a/src/datanode/src/metrics.rs +++ b/src/datanode/src/metrics.rs @@ -17,6 +17,7 @@ use prometheus::*; /// Region request type label. pub const REGION_REQUEST_TYPE: &str = "datanode_region_request_type"; +const REGION_REQUEST_CODE: &str = "code"; pub const REGION_ROLE: &str = "region_role"; pub const REGION_ID: &str = "region_id"; @@ -84,11 +85,12 @@ lazy_static! { ) .unwrap(); - /// Total count of failed insert requests to region server. + /// Total count of failed insert requests to region server, labeled with request + /// type and status code. pub static ref REGION_SERVER_INSERT_FAIL_COUNT: IntCounterVec = register_int_counter_vec!( "greptime_datanode_region_failed_insert_count", "failed region server insert requests count", - &[REGION_REQUEST_TYPE] + &[REGION_REQUEST_TYPE, REGION_REQUEST_CODE] ) .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 10bcf32b13d..527ab0491c0 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -687,7 +687,7 @@ impl RegionServer { } Err(err) => { crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type]) + .with_label_values(&[request_type, err.status_code().as_ref()]) .inc_by(batch_size as u64); Err(err) } @@ -1640,7 +1640,7 @@ impl RegionServerInner { Err(err) => { if matches!(region_change, RegionChange::Ingest) { crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type]) + .with_label_values(&[request_type, err.status_code().as_ref()]) .inc(); } // Removes the region status if the operation fails. From 59eb63c6630612abfc2f011c5e99a8f0c79ff1ed Mon Sep 17 00:00:00 2001 From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Date: Sat, 19 Sep 2026 22:32:09 +0000 Subject: [PATCH 14/63] fix(servers): merge duplicate timeseries across RecordBatches in Prometheus remote read (#9032) * fix(servers): merge timeseries across RecordBatches in Prometheus remote read Signed-off-by: Tyagiquamar * fix(servers): exercise recordbatches_to_timeseries and sort samples with sort_unstable_by_key --------- Signed-off-by: Tyagiquamar Co-authored-by: dennis zhuang --- src/servers/src/prom_store.rs | 174 +++++++++++++++++++++++++++++----- 1 file changed, 151 insertions(+), 23 deletions(-) diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 980504515a8..9182f64ad00 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -362,22 +362,37 @@ pub fn recordbatches_to_timeseries( value_column_name: &str, recordbatches: RecordBatches, ) -> Result> { - Ok(recordbatches - .take() - .into_iter() - .map(|x| recordbatch_to_timeseries(table_name, timestamp_column_name, value_column_name, x)) - .collect::>>()? - .into_iter() - .flatten() - .collect()) + let mut timeseries: Vec = Vec::new(); + let mut timeseries_by_hash: HashMap> = HashMap::new(); + + for recordbatch in recordbatches.take() { + append_recordbatch_to_timeseries( + table_name, + timestamp_column_name, + value_column_name, + recordbatch, + &mut timeseries, + &mut timeseries_by_hash, + )?; + } + + for ts in &mut timeseries { + ts.samples.sort_unstable_by_key(|s| s.timestamp); + } + + timeseries + .sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels)); + Ok(timeseries) } -fn recordbatch_to_timeseries( +fn append_recordbatch_to_timeseries( table: &str, timestamp_column_name: &str, value_column_name: &str, recordbatch: RecordBatch, -) -> Result> { + timeseries: &mut Vec, + timeseries_by_hash: &mut HashMap>, +) -> Result<()> { let ts_column = recordbatch .column_by_name(timestamp_column_name) .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -409,8 +424,6 @@ fn recordbatch_to_timeseries( })?; let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?; - let mut timeseries: Vec = Vec::new(); - let mut timeseries_by_hash: HashMap> = HashMap::new(); let mut previous_timeseries: Option = None; for row in 0..recordbatch.num_rows() { @@ -447,9 +460,7 @@ fn recordbatch_to_timeseries( timeseries[timeseries_index].samples.push(sample); } - timeseries - .sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels)); - Ok(timeseries) + Ok(()) } pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> { @@ -1122,6 +1133,121 @@ mod tests { ); } + #[test] + fn test_recordbatches_to_timeseries_merges_across_batches() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ), + ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), + ])); + + let recordbatches = RecordBatches::try_new( + schema.clone(), + vec![ + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![3000, 1500])) as _, + Arc::new(Float64Vector::from_vec(vec![30.0, 15.0])) as _, + Arc::new(StringVector::from(vec!["host1", "host2"])) as _, + ], + ) + .unwrap(), + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _, + Arc::new(Float64Vector::from_vec(vec![10.0])) as _, + Arc::new(StringVector::from(vec!["host1"])) as _, + ], + ) + .unwrap(), + RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![2000, 2500])) as _, + Arc::new(Float64Vector::from_vec(vec![20.0, 25.0])) as _, + Arc::new(StringVector::from(vec!["host1", "host2"])) as _, + ], + ) + .unwrap(), + ], + ) + .unwrap(); + + let timeseries = recordbatches_to_timeseries( + "cpu_usage", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); + + assert_eq!(2, timeseries.len()); + + assert_eq!( + vec![ + Label { + name: METRIC_NAME_LABEL.to_string(), + value: "cpu_usage".to_string(), + }, + Label { + name: "instance".to_string(), + value: "host1".to_string(), + }, + ], + timeseries[0].labels + ); + assert_eq!( + vec![ + Sample { + value: 10.0, + timestamp: 1000, + }, + Sample { + value: 20.0, + timestamp: 2000, + }, + Sample { + value: 30.0, + timestamp: 3000, + }, + ], + timeseries[0].samples + ); + + assert_eq!( + vec![ + Label { + name: METRIC_NAME_LABEL.to_string(), + value: "cpu_usage".to_string(), + }, + Label { + name: "instance".to_string(), + value: "host2".to_string(), + }, + ], + timeseries[1].labels + ); + assert_eq!( + vec![ + Sample { + value: 15.0, + timestamp: 1500, + }, + Sample { + value: 25.0, + timestamp: 2500, + }, + ], + timeseries[1].samples + ); + } + #[test] fn test_recordbatches_to_timeseries_borrows_and_groups_dictionary_labels() { let arrow_schema = Arc::new(ArrowSchema::new(vec![ @@ -1212,7 +1338,7 @@ mod tests { } #[test] - fn test_recordbatch_to_timeseries_groups_non_contiguous_series() { + fn test_recordbatches_to_timeseries_groups_non_contiguous_series() { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( greptime_timestamp(), @@ -1223,7 +1349,7 @@ mod tests { ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), ])); let recordbatch = RecordBatch::new( - schema, + schema.clone(), vec![ Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _, Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _, @@ -1232,11 +1358,12 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries( + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + let timeseries = recordbatches_to_timeseries( "metric1", greptime_timestamp(), greptime_value(), - recordbatch, + recordbatches, ) .unwrap(); @@ -1259,7 +1386,7 @@ mod tests { } #[test] - fn test_recordbatch_to_timeseries_arrow_label_types_and_nulls() { + fn test_recordbatches_to_timeseries_arrow_label_types_and_nulls() { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( greptime_timestamp(), @@ -1272,7 +1399,7 @@ mod tests { ColumnSchema::new("shard", ConcreteDataType::int32_datatype(), true), ])); let recordbatch = RecordBatch::new( - schema, + schema.clone(), vec![ Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _, Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _, @@ -1289,11 +1416,12 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries( + let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); + let timeseries = recordbatches_to_timeseries( "metric1", greptime_timestamp(), greptime_value(), - recordbatch, + recordbatches, ) .unwrap(); From f29a7d8ff0144257633d5c7777f54e3787d421a6 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 01:59:49 +0000 Subject: [PATCH 15/63] docs(agents): clarify the documentation-update checklist item (#9255) Agents routinely check "This PR requires documentation updates" for changes users never see. Say what the item means: the docs site repo, not rustdoc or in-repo comments. Signed-off-by: Dennis Zhuang --- AGENTS.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 4342279bec9..30e0077caed 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -140,7 +140,12 @@ blast radius requires it. 10. When creating or updating a pull request, follow [`.github/pull_request_template.md`](.github/pull_request_template.md): include the CLA statement, fill the change-intention section with enough detail, and - update checklist items accurately. + update checklist items accurately. "This PR requires documentation updates" + refers to the official docs site repository + ([`GreptimeTeam/docs`](https://github.com/GreptimeTeam/docs)), not to rustdoc + or comments in this repo. Check it only when that site would be wrong or + incomplete without a matching change; a change users cannot observe does not + need it. 11. If the change must also land on release branches, add the matching `backport-` labels (e.g. `backport-v1.3` targets `release/v1.3`). The Backport workflow (`.github/workflows/backport.yml`) cherry-picks a merged PR From 007836ba7afda0ceda1feebea3a3ef08c41fc3fd Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 03:11:06 +0000 Subject: [PATCH 16/63] fix(prometheus): honor label matchers in __name__ values query (#9134) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(prometheus): honor label matchers in __name__ values query `/api/v1/label/__name__/values?match[]={pod="abc"}` dropped every matcher other than `__name__` and returned all metrics in the schema. No error, just the wrong list. Grafana's metrics browser sends this request, so picking a label value there did nothing. Selectors that only constrain `__name__` keep answering from table metadata. A selector constraining an ordinary label now goes to the data: scan each metric engine physical table for distinct `__table_id` in the time range, map the ids back to metric names, then apply the selector's own `__name__` matchers. Only metric engine tables are covered; other engines share no column space to scan. Signed-off-by: Dennis Zhuang * refactor(prometheus): batch-resolve metric names by table id Building a full table-id-to-name map meant walking every table in the schema and holding all of them in memory, just to name the handful the scan returned. Use `tables_by_ids` instead — one batch KV read over the ids the scan actually produced. The catalog walk stays, but only to find the physical tables to scan. Signed-off-by: Dennis Zhuang * fix(promql): read an absent label as the empty string A matcher on a label the series does not carry only worked when the table had no column for it at all. Where the column exists but is NULL on that row -- the norm for logical metrics sharing a metric engine physical table, which holds the union of their label columns -- three-valued logic dropped the row, so `host!="host1"` and `host=""` missed every metric without a host label. Coalesce nullable string label columns to "" for matchers that accept the empty string, rather than only for the OTLP temporality marker. Equality matchers are untouched; they cannot match NULL either way. This is the Prometheus compatibility fix #8970 deliberately kept out of its own scope. The cost is visible in the regex sqlness plan: the predicate becomes a CASE, so the scan loses its LastRow selector and grows a FilterExec. Only negative and empty-accepting matchers pay it. Signed-off-by: Dennis Zhuang * fix(promql): don't panic on pre-epoch label value bounds `rewrite_label_values_query` unwrapped `duration_since(UNIX_EPOCH)`, which returns an error for an instant before the epoch. `start=1969-12-31T23:59:59Z` parses as valid RFC3339, so the request panicked instead of answering. Recover the sign from the error branch, and report a value beyond i64 milliseconds as an error rather than wrapping the cast. Signed-off-by: Dennis Zhuang * refactor(prometheus): drop applicable_matchers, share the distinct scan With the planner reading an absent label as empty, the frontend no longer needs to pre-filter matchers per physical table. Removing that exposed a second problem: a physical table that never took a column from a logical table exposes no `__table_id`, and projecting it failed the whole request. Skip those tables; the only thing that can miss is a metric with no labels. Also pulls out the plan-build-execute-collect sequence the two label value scans had in common. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/error.rs | 13 +- src/frontend/src/instance.rs | 14 + src/frontend/src/instance/promql.rs | 214 +++++++++-- src/query/src/promql/error.rs | 9 + src/query/src/promql/label_values.rs | 55 ++- src/query/src/promql/planner.rs | 7 +- src/query/src/promql/planner/test/delta.rs | 71 +++- src/servers/src/http/prometheus.rs | 336 +++++++++++++++++- src/servers/src/prometheus_handler.rs | 15 + tests-integration/tests/http.rs | 111 ++++++ .../standalone/common/promql/regex.result | 3 +- 11 files changed, 788 insertions(+), 60 deletions(-) diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index c0f9e2fe6bd..9fd2312e736 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -174,6 +174,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Unexpected type {data_type} for column '{column}' of table '{table_name}'"))] + UnexpectedColumnType { + table_name: String, + column: String, + data_type: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to collect recordbatch"))] CollectRecordbatch { #[snafu(implicit)] @@ -414,7 +423,9 @@ impl ErrorExt for Error { Error::RequestQuery { source, .. } => source.status_code(), - Error::CacheRequired { .. } => StatusCode::Internal, + Error::CacheRequired { .. } | Error::UnexpectedColumnType { .. } => { + StatusCode::Internal + } Error::TableNotFound { .. } => StatusCode::TableNotFound, diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 44949f932f5..e16955f2877 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1510,6 +1510,20 @@ impl PrometheusHandler for Instance { .context(ExecuteQuerySnafu) } + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> server_error::Result> { + self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu) + } + async fn query_label_values( &self, metric: String, diff --git a/src/frontend/src/instance/promql.rs b/src/frontend/src/instance/promql.rs index 3df70224703..3ad9ae217ae 100644 --- a/src/frontend/src/instance/promql.rs +++ b/src/frontend/src/instance/promql.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use std::time::SystemTime; @@ -20,20 +21,25 @@ use catalog::information_schema::TABLES; use client::OutputData; use common_catalog::consts::INFORMATION_SCHEMA_NAME; use common_catalog::format_full_table_name; -use common_recordbatch::util; +use common_recordbatch::{RecordBatch, util}; use common_telemetry::tracing; use datafusion_expr::LogicalPlan; +use datatypes::arrow::array::{Array, UInt32Array}; +use futures::StreamExt; use promql_parser::label::{Matcher, Matchers}; use query::promql; use query::promql::planner::PromPlanner; use servers::prometheus; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; +use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; +use store_api::storage::TableId; +use table::TableRef; use crate::error::{ CatalogSnafu, CollectRecordbatchSnafu, ExecLogicalPlanSnafu, PrometheusLabelValuesQueryPlanSnafu, PrometheusMetricNamesQueryPlanSnafu, ReadTableSnafu, - Result, TableNotFoundSnafu, TableSnafu, + Result, TableNotFoundSnafu, TableSnafu, UnexpectedColumnTypeSnafu, }; use crate::instance::Instance; @@ -117,6 +123,173 @@ impl Instance { Ok(results) } + /// Handles a metric names query constrained by matchers on ordinary labels. + /// + /// [`Instance::handle_query_metric_names`] answers from table metadata, which + /// cannot resolve a label matcher: whether a metric carries `pod="abc"` is a + /// property of its data. The metric engine multiplexes the logical tables of + /// one physical table into a region holding the union of their label columns + /// alongside `__table_id`, so a single distinct scan per physical table + /// resolves the matchers for all of its logical tables at once. + /// + /// Only metric engine tables are covered. Tables on other engines share no + /// column space to scan, and a scan per table does not scale to the table + /// counts this API is expected to answer over. + #[tracing::instrument(skip_all)] + pub(crate) async fn handle_query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + let _timer = crate::metrics::PROMQL_QUERY_METRICS_ELAPSED + .with_label_values(&[ctx.get_db_string().as_str()]) + .start_timer(); + + let catalog = ctx.current_catalog(); + let mut table_ids = HashSet::new(); + for physical in self.physical_metric_tables(catalog, schema, ctx).await? { + table_ids.extend( + self.scan_matching_table_ids(physical, &matchers, start, end, ctx) + .await?, + ); + } + + // Batch-resolve only the ids the scan produced. An id dropped between the + // scan and here simply has no entry. + let table_ids = table_ids.into_iter().collect::>(); + let mut names = self + .catalog_manager + .tables_by_ids(catalog, schema, &table_ids) + .await + .context(CatalogSnafu)? + .into_iter() + .map(|table| table.table_info().name.clone()) + .collect::>(); + names.sort_unstable(); + Ok(names) + } + + /// Returns the metric engine physical tables of a schema. + /// + /// Their data regions carry the union of their logical tables' label columns, + /// so scanning these covers every metric of the schema. + async fn physical_metric_tables( + &self, + catalog: &str, + schema: &str, + ctx: &QueryContextRef, + ) -> Result> { + let mut tables = self.catalog_manager.tables(catalog, schema, Some(ctx)); + let mut physical_tables = Vec::new(); + + while let Some(table) = tables.next().await { + let table = table.context(CatalogSnafu)?; + if table.table_info().is_physical_table() { + physical_tables.push(table); + } + } + + Ok(physical_tables) + } + + /// Scans `table` for the distinct values of `column` that match `matchers` + /// within the time range. + /// + /// Callers look the table up and decode the batches; the plan between is the + /// same whether the projected column is a label or `__table_id`. + async fn scan_distinct_column( + &self, + table: TableRef, + matchers: Vec, + column: String, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + let dataframe = self + .query_engine + .read_table(table.clone()) + .with_context(|_| ReadTableSnafu { + table_name: table.table_info().full_table_name(), + })?; + + let scan_plan = dataframe.into_unoptimized_plan(); + let conditions = PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema()) + .context(PrometheusLabelValuesQueryPlanSnafu)?; + let logical_plan = promql::label_values::rewrite_label_values_query( + table, scan_plan, conditions, column, start, end, + ) + .context(PrometheusLabelValuesQueryPlanSnafu)?; + + let results = self + .query_engine + .execute(logical_plan, ctx.clone()) + .await + .context(ExecLogicalPlanSnafu)?; + + match results.data { + OutputData::Stream(stream) => { + util::collect(stream).await.context(CollectRecordbatchSnafu) + } + OutputData::RecordBatches(rbs) => Ok(rbs.take()), + _ => unreachable!("should not happen"), + } + } + + /// Returns the `__table_id`s of `physical` carrying a row that matches every + /// matcher within the time range. + async fn scan_matching_table_ids( + &self, + physical: TableRef, + matchers: &[Matcher], + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result> { + // `__table_id` attributes a row to its logical table, and a physical + // table that never took a column from one does not expose it. Skipping + // such a table can only miss a metric carrying no labels at all. + if physical + .schema() + .column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + .is_none() + { + return Ok(Vec::new()); + } + + let table_name = physical.table_info().full_table_name(); + let batches = self + .scan_distinct_column( + physical, + matchers.to_vec(), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + start, + end, + ctx, + ) + .await?; + + let mut table_ids = Vec::new(); + for batch in batches { + // Only one column in results, ensured by `rewrite_label_values_query`. + let column = batch.column(0); + let ids = column + .as_any() + .downcast_ref::() + .with_context(|| UnexpectedColumnTypeSnafu { + table_name: table_name.clone(), + column: DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + data_type: column.data_type().to_string(), + })?; + table_ids.extend(ids.iter().flatten()); + } + + Ok(table_ids) + } + /// Handles label values query request, returns the values. #[tracing::instrument(skip_all)] pub(crate) async fn handle_query_label_values( @@ -150,40 +323,9 @@ impl Instance { .context(TableSnafu); } - let dataframe = self - .query_engine - .read_table(table.clone()) - .with_context(|_| ReadTableSnafu { - table_name: full_table_name, - })?; - - let scan_plan = dataframe.into_unoptimized_plan(); - let filter_conditions = - PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema()) - .context(PrometheusLabelValuesQueryPlanSnafu)?; - let logical_plan = promql::label_values::rewrite_label_values_query( - table, - scan_plan, - filter_conditions, - label_name, - start, - end, - ) - .context(PrometheusLabelValuesQueryPlanSnafu)?; - - let results = self - .query_engine - .execute(logical_plan, ctx.clone()) - .await - .context(ExecLogicalPlanSnafu)?; - - let batches = match results.data { - OutputData::Stream(stream) => util::collect(stream) - .await - .context(CollectRecordbatchSnafu)?, - OutputData::RecordBatches(rbs) => rbs.take(), - _ => unreachable!("should not happen"), - }; + let batches = self + .scan_distinct_column(table, matchers, label_name, start, end, ctx) + .await?; let mut results = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum()); for batch in batches { diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index dba55c8df8b..45e08349623 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::any::Any; +use std::time::SystemTime; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -202,6 +203,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Time out of the representable millisecond range: {:?}", time))] + SystemTimeOutOfRange { + time: SystemTime, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("vector cannot contain metrics with the same labelset"))] SameLabelSet { #[snafu(implicit)] @@ -243,6 +251,7 @@ impl ErrorExt for Error { | UnsupportedMatcherOp { .. } | SameLabelSet { .. } | TimestampOutOfRange { .. } + | SystemTimeOutOfRange { .. } | InvalidRegularExpression { .. } | InvalidDestinationLabelName { .. } => StatusCode::InvalidArguments, diff --git a/src/query/src/promql/label_values.rs b/src/query/src/promql/label_values.rs index 4f6ea42a60b..44c173477bb 100644 --- a/src/query/src/promql/label_values.rs +++ b/src/query/src/promql/label_values.rs @@ -23,9 +23,31 @@ use snafu::{OptionExt, ResultExt}; use table::TableRef; use crate::promql::error::{ - DataFusionPlanningSnafu, Result, TimeIndexNotFoundSnafu, TimestampOutOfRangeSnafu, + DataFusionPlanningSnafu, Result, SystemTimeOutOfRangeSnafu, TimeIndexNotFoundSnafu, + TimestampOutOfRangeSnafu, }; +/// Converts a [`SystemTime`] to a millisecond [`Timestamp`]. +/// +/// `duration_since` reports an instant before the epoch as an error rather than +/// a negative duration, and an RFC3339 query parameter can name one, so the sign +/// is recovered here instead of unwrapping. +fn millis_since_epoch(time: SystemTime) -> Result { + let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) { + Ok(duration) => (duration.as_millis(), false), + Err(earlier) => (earlier.duration().as_millis(), true), + }; + let millis = i64::try_from(millis) + .ok() + .with_context(|| SystemTimeOutOfRangeSnafu { time })?; + + Ok(Timestamp::new_millisecond(if before_epoch { + -millis + } else { + millis + })) +} + fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr { time_index_expr .clone() @@ -67,14 +89,12 @@ pub fn rewrite_label_values_query( })?; // We only support millisecond precision at most. - let start = - Timestamp::new_millisecond(start.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64); + let start = millis_since_epoch(start)?; let start = start.convert_to(unit).context(TimestampOutOfRangeSnafu { timestamp: start.value(), unit, })?; - let end = - Timestamp::new_millisecond(end.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64); + let end = millis_since_epoch(end)?; let end = end.convert_to(unit).context(TimestampOutOfRangeSnafu { timestamp: end.value(), unit, @@ -98,3 +118,28 @@ pub fn rewrite_label_values_query( Ok(logical_plan) } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + #[test] + fn millis_before_the_epoch_are_negative() { + // `SystemTime::duration_since` reports these as an error; unwrapping it + // panicked on any request whose RFC3339 start named a pre-epoch instant. + let time = UNIX_EPOCH - Duration::from_millis(1); + assert_eq!(millis_since_epoch(time).unwrap().value(), -1); + assert_eq!(millis_since_epoch(UNIX_EPOCH).unwrap().value(), 0); + + let time = UNIX_EPOCH + Duration::from_millis(1); + assert_eq!(millis_since_epoch(time).unwrap().value(), 1); + } + + #[test] + fn millis_beyond_i64_are_rejected() { + let time = UNIX_EPOCH + Duration::from_secs(1 << 60); + assert!(millis_since_epoch(time).is_err()); + } +} diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 17cc36a41da..bea2cae586c 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -2687,13 +2687,17 @@ impl PromPlanner { let accepts_empty = matcher.is_match(""); let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str()); + // Prometheus reads a label a series does not carry as the empty + // string. A row can miss a label two ways: the table has no column + // for it, or the column exists but is NULL on that row — the latter + // is the norm for logical metrics sharing a physical table, which + // holds the union of their label columns. let col = if let Some(column_name) = column_name { let column = DfExpr::Column(Column::from_name(&column_name)); let field = table_schema .index_of_column_by_name(None, &column_name) .map(|index| table_schema.field(index)); if accepts_empty - && column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL && let Some(data_type) = field .filter(|field| { field.is_nullable() @@ -2712,7 +2716,6 @@ impl PromPlanner { } } else { DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) - .alias(matcher.name.clone()) }; let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None); let expr = match matcher.op { diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index f6c37bb24a1..3727ea9a30c 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -524,7 +524,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { } #[tokio::test] -async fn temporality_matchers_treat_null_as_absent() { +async fn matchers_read_absent_labels_as_empty() { let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL; let schema = Arc::new(ArrowSchema::new(vec![Field::new( marker, @@ -596,31 +596,82 @@ async fn temporality_matchers_treat_null_as_absent() { ); } - let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( + // The rule is about NULL, not about the marker: any nullable label column + // can be NULL where the series does not carry the label. A non-nullable one + // has nothing to normalize. + for (nullable, wants_coalesce) in [(true, true), (false, false)] { + let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "label", + ArrowDataType::Utf8, + nullable, + )])); + let ordinary_scan = LogicalPlanBuilder::scan( + "ordinary_labels", + provider_as_source(Arc::new( + MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(), + )), + None, + ) + .unwrap() + .build() + .unwrap(); + let PromExpr::VectorSelector(selector) = + parser::parse(r#"metric{label!="delta"}"#).unwrap() + else { + unreachable!() + }; + let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema()) + .unwrap() + .iter() + .map(ToString::to_string) + .join(" AND "); + assert_eq!( + wants_coalesce, + expressions.contains("coalesce"), + "nullable={nullable}: {expressions}" + ); + } + + // The other way a label goes absent is having no column for it at all. The + // literal standing in for the label has to be usable as a predicate. + let schema = Arc::new(ArrowSchema::new(vec![Field::new( "label", ArrowDataType::Utf8, true, )])); - let ordinary_scan = LogicalPlanBuilder::scan( + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(vec![Some("value")]))], + ) + .unwrap(); + let scan = LogicalPlanBuilder::scan( "ordinary_labels", provider_as_source(Arc::new( - MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(), + MemTable::try_new(schema, vec![vec![batch]]).unwrap(), )), None, ) .unwrap() .build() .unwrap(); - let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{label!="delta"}"#).unwrap() + let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{absent!="delta"}"#).unwrap() else { unreachable!() }; - let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema()) + let expressions = PromPlanner::matchers_to_expr(selector.matchers, scan.schema()).unwrap(); + assert_eq!( + r#"Utf8("") != Utf8("delta")"#, + expressions.iter().map(ToString::to_string).join(" AND ") + ); + + let plan = LogicalPlanBuilder::from(scan) + .filter(conjunction(expressions).unwrap()) .unwrap() - .iter() - .map(ToString::to_string) - .join(" AND "); - assert!(!expressions.contains("coalesce"), "{expressions}"); + .build() + .unwrap(); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let rows = batches.iter().map(|batch| batch.num_rows()).sum::(); + assert_eq!(1, rows); } #[tokio::test] diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 1f9988ede4c..420a0c8f4b4 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -55,7 +55,7 @@ use promql_parser::parser::{ AggregateExpr, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorSelector, }; -use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryStatement}; +use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement}; use serde::de::{self, MapAccess, Visitor}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -74,7 +74,7 @@ use table::requests::{ pub use super::result::prometheus_resp::{PromSampleValue, PrometheusJsonResponse}; use crate::error::{ CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu, - NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, + NotSupportedSnafu, ParseTimestampSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, }; use crate::http::header::collect_plan_metrics; use crate::otlp::metrics::ucum_to_openmetrics_unit; @@ -1711,9 +1711,35 @@ pub async fn label_values_query( ); let catalog_manager = handler.catalog_manager(); - let mut table_names = try_call_return_response!( - retrieve_table_names(&query_ctx, catalog_manager, matches).await - ); + // An empty `match[]` enumerates every metric; otherwise only the + // selectors answerable from metadata go down that path. + let enumerate_all = matches.is_empty(); + let (metadata_selectors, label_selectors) = + try_call_return_response!(split_selectors_by_label_use(&matches)); + + let mut table_names = if enumerate_all || !metadata_selectors.is_empty() { + try_call_return_response!( + retrieve_table_names(&query_ctx, catalog_manager, metadata_selectors).await + ) + } else { + Vec::new() + }; + + if !label_selectors.is_empty() { + table_names.extend(try_call_return_response!( + retrieve_table_names_by_labels( + &handler, + label_selectors, + params.start.as_deref(), + params.end.as_deref(), + &query_ctx, + ) + .await + )); + table_names.sort_unstable(); + table_names.dedup(); + } + table_names = try_call_return_response!( handler .filter_metadata_metric_names( @@ -1897,6 +1923,131 @@ fn take_metric_name(selector: &mut VectorSelector) -> Option { Some(name) } +/// Removes every `__name__` matcher from the selector and returns them, so the +/// rest can be planned as column predicates. A name given as `VectorSelector::name` +/// comes back as an equality matcher, making both spellings filter alike. +fn take_metric_name_matchers(selector: &mut VectorSelector) -> Vec { + let mut taken = Vec::new(); + if let Some(name) = selector.name.take() { + taken.push(Matcher::new(MatchOp::Equal, METRIC_NAME_LABEL, &name)); + } + + let (name_matchers, rest) = std::mem::take(&mut selector.matchers.matchers) + .into_iter() + .partition(|matcher| matcher.name == METRIC_NAME_LABEL); + selector.matchers.matchers = rest; + taken.extend(name_matchers); + + taken +} + +/// Whether a metric name satisfies every `__name__` matcher of one selector. +/// +/// Negated matchers are honoured here, unlike in [`retrieve_table_names`] where +/// they keep every table so the caller authorizes the full candidate set: the +/// names reaching this point are already narrowed by the data, so filtering them +/// can only remove names, never widen what a caller gets to see. +fn metric_name_matches(table_name: &str, matchers: &[Matcher]) -> bool { + matchers.iter().all(|matcher| match &matcher.op { + MatchOp::Equal => table_name == matcher.value, + MatchOp::NotEqual => table_name != matcher.value, + MatchOp::Re(re) => re.is_match(table_name), + MatchOp::NotRe(re) => !re.is_match(table_name), + }) +} + +/// Whether a matcher constrains an ordinary label. The others name the metric, +/// the database or the field, none of which is a column to scan. +fn is_ordinary_label_matcher(matcher: &Matcher) -> bool { + matcher.name != METRIC_NAME_LABEL + && matcher.name != FIELD_NAME_LABEL + && !is_database_selection_label(&matcher.name) +} + +/// Splits `match[]` selectors by whether they constrain an ordinary label. The +/// first group is answerable from table metadata; the second needs the data read +/// and is returned as parsed selectors. +/// +/// `or` matchers stay in the metadata group, which ignores them, rather than +/// being silently dropped from a data scan that cannot express them. +fn split_selectors_by_label_use(matches: &[String]) -> Result<(Vec, Vec)> { + let mut metadata_only = Vec::new(); + let mut with_labels = Vec::new(); + + for selector in matches { + let expr = promql_parser::parser::parse(selector) + .map_err(|reason| InvalidQuerySnafu { reason }.build())?; + let PromqlExpr::VectorSelector(vector_selector) = expr else { + return InvalidQuerySnafu { + reason: "expected vector selector".to_string(), + } + .fail(); + }; + + let constrains_labels = vector_selector.matchers.or_matchers.is_empty() + && vector_selector + .matchers + .matchers + .iter() + .any(is_ordinary_label_matcher); + if constrains_labels { + with_labels.push(vector_selector); + } else { + metadata_only.push(selector.clone()); + } + } + + Ok((metadata_only, with_labels)) +} + +/// Resolves selectors constraining ordinary labels into metric names: the data +/// answers the label matchers, then each selector's `__name__` matchers narrow +/// the names it found. +async fn retrieve_table_names_by_labels( + handler: &PrometheusHandlerRef, + selectors: Vec, + start: Option<&str>, + end: Option<&str>, + query_ctx: &QueryContextRef, +) -> Result> { + let start_arg = start.map(str::to_string).unwrap_or_else(yesterday_rfc3339); + let end_arg = end.map(str::to_string).unwrap_or_else(current_time_rfc3339); + let start = QueryLanguageParser::parse_promql_timestamp(&start_arg).with_context(|_| { + ParseTimestampSnafu { + timestamp: start_arg.clone(), + } + })?; + let end = QueryLanguageParser::parse_promql_timestamp(&end_arg).with_context(|_| { + ParseTimestampSnafu { + timestamp: end_arg.clone(), + } + })?; + + let schema = query_ctx.current_schema(); + let mut table_names = Vec::new(); + for mut selector in selectors { + let name_matchers = take_metric_name_matchers(&mut selector); + // The database and field matchers name no column, and the metadata path + // ignores them too. + let label_matchers = selector + .matchers + .matchers + .into_iter() + .filter(is_ordinary_label_matcher) + .collect(); + let matched = handler + .query_metric_names_by_labels(label_matchers, &schema, start, end, query_ctx) + .await?; + table_names.extend( + matched + .into_iter() + .filter(|name| metric_name_matches(name, &name_matchers)), + ); + } + + Ok(table_names) +} + async fn retrieve_table_names( query_ctx: &QueryContext, catalog_manager: CatalogManagerRef, @@ -2407,6 +2558,10 @@ mod tests { deny_operation: bool, denied_table: Option<&'static str>, metric_names: Vec, + /// Names the label-matcher path resolves, kept apart from `metric_names` + /// so a test can tell which path answered. + label_metric_names: Vec, + label_lookups: Mutex>>, queries: Mutex>, ordered_outputs: Mutex>, } @@ -2482,6 +2637,18 @@ mod tests { Ok(self.metric_names.clone()) } + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + _: &str, + _: std::time::SystemTime, + _: std::time::SystemTime, + _: &QueryContextRef, + ) -> Result> { + self.label_lookups.lock().unwrap().push(matchers); + Ok(self.label_metric_names.clone()) + } + async fn query_label_values( &self, _: String, @@ -2547,6 +2714,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -2600,6 +2769,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -2684,6 +2855,8 @@ mod tests { deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), })), @@ -2700,6 +2873,153 @@ mod tests { ); } + /// A handler over the logical metric tables `cpu_user` and `cpu_system`, + /// which is what the metadata path enumerates, with the label-matcher path + /// answering `label_metric_names`. + fn label_values_handler(label_metric_names: Vec<&str>) -> Arc { + let mut cpu_user = test_table_info( + 1024, + "cpu_user", + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + Arc::new(Schema::new(vec![])), + ); + cpu_user.meta.options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "physical_metrics".to_string(), + ); + let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&cpu_user)); + let mut cpu_system = cpu_user.clone(); + cpu_system.ident.table_id = 1025; + cpu_system.name = "cpu_system".to_string(); + manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: cpu_system.name.clone(), + table_id: cpu_system.table_id(), + table: EmptyTable::from_table_info(&cpu_system), + }) + .unwrap(); + + Arc::new(TestPrometheusHandler { + catalog_manager: manager, + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + label_metric_names: label_metric_names.into_iter().map(String::from).collect(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }) + } + + async fn query_metric_name_values( + handler: Arc, + matches: Vec<&str>, + ) -> Vec { + let state: PrometheusHandlerRef = handler; + let response = label_values_query( + State(state), + Path(METRIC_NAME_LABEL.to_string()), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Query(LabelValueQuery { + matches: Matches(matches.into_iter().map(String::from).collect()), + ..Default::default() + }), + ) + .await; + + assert!( + response.status_code.is_none(), + "status={:?}, error={:?}", + response.status_code, + response.error + ); + match response.data { + PrometheusResponse::LabelValues(values) => values, + other => panic!("expected label values, got {other:?}"), + } + } + + #[tokio::test] + async fn label_matchers_resolve_metric_names_from_data() { + let handler = label_values_handler(vec!["cpu_user"]); + let values = query_metric_name_values(handler.clone(), vec![r#"{pod="abc"}"#]).await; + + // The metadata path would have enumerated both metrics. + assert_eq!(vec!["cpu_user".to_string()], values); + + let lookups = handler.label_lookups.lock().unwrap(); + assert_eq!(1, lookups.len()); + assert_eq!( + vec!["pod".to_string()], + lookups[0] + .iter() + .map(|matcher| matcher.name.clone()) + .collect::>() + ); + } + + #[tokio::test] + async fn metric_name_matchers_narrow_data_resolved_names() { + let handler = label_values_handler(vec!["cpu_user", "cpu_system"]); + let values = + query_metric_name_values(handler, vec![r#"{__name__=~"cpu_u.*", pod="abc"}"#]).await; + + assert_eq!(vec!["cpu_user".to_string()], values); + } + + #[tokio::test] + async fn special_matchers_are_stripped_before_the_data_lookup() { + let handler = label_values_handler(vec!["cpu_user"]); + let values = query_metric_name_values( + handler.clone(), + vec![r#"{pod="abc", __field__="value", __database__="public"}"#], + ) + .await; + + assert_eq!(vec!["cpu_user".to_string()], values); + let lookups = handler.label_lookups.lock().unwrap(); + assert_eq!( + vec!["pod".to_string()], + lookups[0] + .iter() + .map(|matcher| matcher.name.clone()) + .collect::>() + ); + } + + #[tokio::test] + async fn database_and_field_matchers_stay_on_the_metadata_path() { + let handler = label_values_handler(vec!["never_returned"]); + let values = query_metric_name_values( + handler.clone(), + vec![r#"{__name__=~"cpu_.*", __field__="value", __database__="other"}"#], + ) + .await; + + assert_eq!( + vec!["cpu_system".to_string(), "cpu_user".to_string()], + values + ); + assert!(handler.label_lookups.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn empty_match_still_enumerates_every_metric() { + let handler = label_values_handler(Vec::new()); + let values = query_metric_name_values(handler, Vec::new()).await; + + assert_eq!( + vec!["cpu_system".to_string(), "cpu_user".to_string()], + values + ); + } + #[tokio::test] async fn test_series_query_expands_metric_name_regex() { let cpu_user = test_table_info( @@ -2728,6 +3048,8 @@ mod tests { deny_operation: false, denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); @@ -3635,6 +3957,8 @@ mod tests { deny_operation: true, denied_table: None, metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), })), @@ -3649,6 +3973,8 @@ mod tests { deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), queries: Mutex::new(Vec::new()), ordered_outputs: Mutex::new(Vec::new()), }); diff --git a/src/servers/src/prometheus_handler.rs b/src/servers/src/prometheus_handler.rs index 4f4b7c5d19c..1e211b501ab 100644 --- a/src/servers/src/prometheus_handler.rs +++ b/src/servers/src/prometheus_handler.rs @@ -184,6 +184,21 @@ pub trait PrometheusHandler { ctx: &QueryContextRef, ) -> Result>; + /// Query metric table names that carry data matching `matchers` in the time + /// range. `matchers` must hold only ordinary label matchers: `__name__` + /// names a table and the database and field matchers name no column, so the + /// caller resolves all three before calling. + /// + /// Only metric engine tables are covered. + async fn query_metric_names_by_labels( + &self, + matchers: Vec, + schema: &str, + start: SystemTime, + end: SystemTime, + ctx: &QueryContextRef, + ) -> Result>; + async fn query_label_values( &self, metric: String, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 12b9ac883d1..bd761f39a31 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1524,6 +1524,117 @@ pub async fn test_prom_http_api(store_type: StorageType) { .unwrap() ); + // query `__name__` by a matcher on an ordinary label: the metric engine + // physical tables are scanned, so only metrics carrying the label value are + // returned. `demo_metrics` shares `phy` with `demo` but has no `host` value. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=0&end=600") + .send() + .await; + let status = res.status(); + let text = res.text().await; + assert_eq!(status, StatusCode::OK, "{text}"); + let prom_resp = serde_json::from_str::(&text).unwrap(); + assert_eq!(prom_resp.status, "success"); + assert!(prom_resp.error.is_none()); + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + + // `__name__` matchers narrow the names the data resolved: `multi_labels` + // also carries `idc="idc1"` but its name does not match. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\"demo.*\", idc=\"idc1\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!(prom_resp.status, "success"); + assert!(prom_resp.error.is_none()); + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo_metrics", + "demo_metrics_with_nanos", + ])) + .unwrap() + ); + + // The time range selects the series: `demo` carries `host="host2"` only at + // t=600, so narrowing the range drops it while `multi_labels` at t=0 stays. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=300") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["multi_labels"])).unwrap() + ); + + // Logical metrics sharing a physical table have NULL in the label columns + // they don't use. Prometheus reads a label a series doesn't carry as the + // empty string, so `demo_metrics` and `demo_metrics_with_nanos` — neither of + // which has a `host` label — match both of these. + // + // `.%2B` is `.+`; a bare `+` decodes to a space in a query string. Grafana + // encodes it the same way. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host=\"\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo_metrics", + "demo_metrics_with_nanos", + ])) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host!=\"host1\"}&start=0&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!([ + "demo", + "demo_metrics", + "demo_metrics_with_nanos", + "multi_labels", + ])) + .unwrap() + ); + + // A pre-epoch RFC3339 bound is a valid range, not a panic. + let res = client + .get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=1969-12-31T23:59:59Z&end=600") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let prom_resp = res.json::().await; + assert_eq!( + prom_resp.data, + serde_json::from_value::(json!(["demo", "multi_labels"])).unwrap() + ); + // buildinfo let res = client .get("/v1/prometheus/api/v1/status/buildinfo") diff --git a/tests/cases/standalone/common/promql/regex.result b/tests/cases/standalone/common/promql/regex.result index 7e0132c2f40..c76b743eafa 100644 --- a/tests/cases/standalone/common/promql/regex.result +++ b/tests/cases/standalone/common/promql/regex.result @@ -146,8 +146,9 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED +|_|_|_FilterExec: CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END = REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["CASE WHEN host IS NOT NULL THEN host ELSE Dictionary(UInt32, Utf8(\"\")) END = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ From e99f75f9b3a0911074da6f3ac66a78b51a63a31d Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 03:12:51 +0000 Subject: [PATCH 17/63] fix(mito2): hold the bulk compact permit for the whole blocking task (#9252) CompactDispatcher acquired a borrowed permit inside the async wrapper and dropped it as soon as the blocking task was submitted, so the semaphore never bounded the compactions that actually ran. Acquire an owned permit and move it into the blocking closure so it covers the queue wait and the merge itself. Signed-off-by: Dennis Zhuang --- src/mito2/src/memtable/bulk.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index ea2a1225f13..94ccac472df 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1545,11 +1545,15 @@ impl CompactDispatcher { fn dispatch_compact(&self, task: MemCompactTask) { let semaphore = self.semaphore.clone(); common_runtime::spawn_global(async move { - let Ok(_permit) = semaphore.acquire().await else { + let Ok(permit) = semaphore.acquire_owned().await else { return; }; common_runtime::spawn_blocking_global(move || { + // The async task above returns as soon as the blocking task is + // submitted, so the permit has to travel with it to bound the + // compaction that actually runs. + let _permit = permit; if let Err(e) = task.compact() { common_telemetry::error!(e; "Failed to compact memtable, region: {}", task.metadata.region_id); } From 798adb64e7474e7040f8d5996af22c4cca9091a6 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Sun, 20 Sep 2026 03:13:52 +0000 Subject: [PATCH 18/63] feat(mito2): make range index reads and builds opt-in (#9219) * feat: add config flag to control range index reads Signed-off-by: evenyag * fix: disable range index builds when configured and default to off Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 2 + config/datanode.example.toml | 4 + config/standalone.example.toml | 4 + src/mito2/src/config.rs | 20 ++++ src/mito2/src/engine.rs | 1 + src/mito2/src/engine/scan_test.rs | 51 +++++++- src/mito2/src/read/scan_region.rs | 27 ++++- src/mito2/src/series_index/maintenance.rs | 97 ++++++++------- src/mito2/src/series_index/task.rs | 5 + src/mito2/src/series_index/tests.rs | 140 +++++++++++++++++++++- src/mito2/src/worker.rs | 1 + tests-integration/tests/http.rs | 1 + 12 files changed, 302 insertions(+), 51 deletions(-) diff --git a/config/config.md b/config/config.md index 4dcffe40654..651cbec6e85 100644 --- a/config/config.md +++ b/config/config.md @@ -177,6 +177,7 @@ | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | | `region_engine.mito.max_background_flushes` | Integer | Auto | Max number of running background flush jobs (default: 1/2 of cpu cores). | @@ -638,6 +639,7 @@ | `region_engine.mito.experimental_manifest_keep_removed_file_ttl` | String | `1h` | How long to keep removed files in the `removed_files` field of manifest
after they are removed from manifest.
files will only be removed from `removed_files` field
if both `keep_removed_file_count` and `keep_removed_file_ttl` is reached. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | | `region_engine.mito.max_background_flushes` | Integer | Auto | Max number of running background flush jobs (default: 1/2 of cpu cores). | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 814c78c3ade..f99e1e124f9 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -501,6 +501,10 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Whether to build and query range indexes when series indexes are enabled. +## Obsolete range-index metadata and files are still cleaned up when disabled. +#+ experimental_enable_range_index = false + ## Interval between series-index maintenance runs. Zero uses the default of 5 min. #+ experimental_series_index_maintenance_interval = "5m" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 54a1555232e..57893881146 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -694,6 +694,10 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Whether to build and query range indexes when series indexes are enabled. +## Obsolete range-index metadata and files are still cleaned up when disabled. +#+ experimental_enable_range_index = false + ## Interval between series-index maintenance runs. Zero uses the default of 5 min. #+ experimental_series_index_maintenance_interval = "5m" diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index f525757a268..8ca4326e2b6 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -92,6 +92,9 @@ pub struct MitoConfig { /// Under development; do not enable. Whether to enable series indexes (default false). /// Indexes are stored on the local filesystem under `{data_home}/series_index`. pub experimental_enable_series_index: bool, + /// Whether to build and query range indexes when series indexes are enabled (default false). + /// Obsolete range-index metadata and files are still cleaned up when disabled. + pub experimental_enable_range_index: bool, /// Interval between series-index maintenance runs (default 5 min). Zero uses the default. #[serde(with = "humantime_serde")] pub experimental_series_index_maintenance_interval: Duration, @@ -217,6 +220,7 @@ impl Default for MitoConfig { compress_manifest: false, max_background_index_builds: divide_num_cpus(8), experimental_enable_series_index: false, + experimental_enable_range_index: false, experimental_series_index_maintenance_interval: DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL, experimental_series_index_bucket_width: Duration::from_secs(5 * 24 * 60 * 60), @@ -429,18 +433,21 @@ mod tests { #[test] fn test_series_index_config() { assert!(!MitoConfig::default().experimental_enable_series_index); + assert!(!MitoConfig::default().experimental_enable_range_index); assert_eq!( MitoConfig::default().experimental_series_index_bucket_width, Duration::from_secs(5 * 24 * 60 * 60) ); let mut config: MitoConfig = toml::from_str( "experimental_enable_series_index = true + experimental_enable_range_index = false experimental_series_index_maintenance_interval = '30s' experimental_series_index_bucket_width = '2days'", ) .unwrap(); config.sanitize("/data").unwrap(); assert!(config.experimental_enable_series_index); + assert!(!config.experimental_enable_range_index); assert_eq!( config.experimental_series_index_maintenance_interval, Duration::from_secs(30) @@ -461,6 +468,19 @@ mod tests { ); } + #[test] + fn test_range_index_config() { + let config: MitoConfig = toml::from_str("experimental_enable_series_index = true").unwrap(); + assert!(!config.experimental_enable_range_index); + for enabled in [false, true] { + let config: MitoConfig = + toml::from_str(&format!("experimental_enable_range_index = {enabled}")).unwrap(); + assert_eq!(enabled, config.experimental_enable_range_index); + let restored: MitoConfig = toml::from_str(&toml::to_string(&config).unwrap()).unwrap(); + assert_eq!(config, restored); + } + } + #[test] fn test_experimental_series_scan_v2_config() { assert!(MitoConfig::default().experimental_series_scan_v2); diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index e1458297a16..fbfd339cff7 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -1156,6 +1156,7 @@ impl EngineInner { CacheStrategy::EnableAll(cache_manager), ) .with_series_index(series_index) + .with_ignore_range_index(!self.config.experimental_enable_range_index) .with_query_stat_counters(region.region_stats.query_stat_counters()) .with_max_concurrent_scan_files(self.config.max_concurrent_scan_files) .with_scan_memory_pool(self.scan_memory_pool.clone()) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index e5362aa4ddf..bf089737f35 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1234,17 +1234,23 @@ async fn test_series_scan_with_format(flat_format: bool) { async fn test_two_phase_series_scan() { for use_index in [false, true] { for use_range_index in [false, true] { - check_two_phase_series_scan(use_index, use_range_index).await; + check_two_phase_series_scan(use_index, use_range_index, true).await; } + check_two_phase_series_scan(use_index, true, false).await; } } -async fn check_two_phase_series_scan(use_index: bool, use_range_index: bool) { +async fn check_two_phase_series_scan( + use_index: bool, + use_range_index: bool, + enable_range_index: bool, +) { let mut env = TestEnv::with_prefix("test_two_phase_series_scan").await; let engine = env .create_engine(MitoConfig { experimental_series_scan_v2: true, experimental_enable_series_index: true, + experimental_enable_range_index: enable_range_index, ..Default::default() }) .await; @@ -1597,8 +1603,8 @@ async fn check_two_phase_series_scan(use_index: bool, use_range_index: bool) { ); if use_range_index { - // A fresh query must read the cataloged index: missing files must not - // silently switch back to the primary-key path. + // Enabled queries must read the cataloged index; disabled queries must + // succeed even when that file is unavailable. let region = engine.find_region(region_id).unwrap(); let version = region.series_index_version_control.current(); let file_id = *version.range_indexes.iter().next().unwrap(); @@ -1621,8 +1627,41 @@ async fn check_two_phase_series_scan(use_index: bool, use_range_index: bool) { ) .await .unwrap(); - if let Ok(stream) = scanner.scan().await { - assert!(stream.try_collect::>().await.is_err()); + if enable_range_index { + if let Ok(stream) = scanner.scan().await { + assert!(stream.try_collect::>().await.is_err()); + } + } else { + let batches = scanner + .scan() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in batches { + let tags = batch.column_by_name("tag_0").unwrap(); + let fields = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + for row in 0..batch.num_rows() { + rows.push(( + datatypes::arrow_array::string_array_value_at_index(tags, row) + .unwrap() + .to_string(), + fields.value(row), + timestamps.value(row), + )); + } + } + rows.sort(); + assert_eq!(actual_rows, rows); } } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 13864b05c24..7c93d137c04 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -250,6 +250,8 @@ pub(crate) struct ScanRegion { scan_memory_pool: Arc, /// Whether to enable the experimental two-phase metric series scan. experimental_series_scan_v2: bool, + /// Whether to ignore range indexes during scans. + ignore_range_index: bool, /// Whether to ignore inverted index. ignore_inverted_index: bool, /// Whether to ignore fulltext index. @@ -286,6 +288,7 @@ impl ScanRegion { max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), experimental_series_scan_v2: false, + ignore_range_index: false, ignore_inverted_index: false, ignore_fulltext_index: false, ignore_bloom_filter: false, @@ -339,6 +342,13 @@ impl ScanRegion { self } + /// Sets whether to ignore range indexes during scans. + #[must_use] + pub(crate) fn with_ignore_range_index(mut self, ignore: bool) -> Self { + self.ignore_range_index = ignore; + self + } + /// Sets whether to ignore inverted index. #[must_use] pub(crate) fn with_ignore_inverted_index(mut self, ignore: bool) -> Self { @@ -574,6 +584,7 @@ impl ScanRegion { let input = ScanInput::builder(self.access_layer, mapper) .with_series_index(self.series_index) + .with_ignore_range_index(self.ignore_range_index) .with_time_range(Some(time_range)) .with_predicate(predicate) .with_memtables(mem_range_builders) @@ -944,6 +955,8 @@ fn time_range_covers_file(time_range: Option<&TimestampRange>, file: &FileHandle pub struct ScanInput { /// Pinned series-index snapshot and its storage, when configured. pub(crate) series_index: Option, + /// Whether to ignore range indexes while retaining series-index candidate discovery. + ignore_range_index: bool, /// Region SST access layer. access_layer: AccessLayerRef, /// Maps projected Batches to RecordBatches. @@ -1043,6 +1056,7 @@ impl ScanInput { ScanInputBuilder { input: ScanInput { series_index: None, + ignore_range_index: false, access_layer, read_cols: mapper.read_columns().clone(), mapper: Arc::new(mapper), @@ -1107,6 +1121,13 @@ impl ScanInput { } impl ScanInputBuilder { + /// Sets whether to ignore range indexes during scans. + #[must_use] + pub(crate) fn with_ignore_range_index(mut self, ignore: bool) -> Self { + self.input.ignore_range_index = ignore; + self + } + /// Sets the pinned series-index context for candidate discovery. #[must_use] pub(crate) fn with_series_index( @@ -1577,7 +1598,11 @@ impl ScanInput { let reader = self .access_layer .read_sst(file.clone()) - .series_index(self.series_index.clone()) + .series_index( + self.series_index + .clone() + .filter(|_| !self.ignore_range_index), + ) .predicate(predicate) .projection(Some(self.read_cols.clone())) .json2_rewrite_targets(self.json2_rewrite_targets.clone()) diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index 4cf8ade65c0..d782cf5a2c1 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -81,6 +81,7 @@ pub(crate) async fn reconcile_series_indexes( requested_bucket_width: Duration, now_ms: i64, purger: IndexFilePurger, + enable_range_index: bool, ) -> Result { let total_start = Instant::now(); // Use this snapshot throughout reconciliation, even if the region version advances. @@ -102,15 +103,16 @@ pub(crate) async fn reconcile_series_indexes( now_ms, &purger, &mut unpublished, + enable_range_index, ) .await?; SERIES_INDEX_RECONCILE_ELAPSED .with_label_values(&["build"]) .observe(build_start.elapsed().as_secs_f64()); - // Persist both catalogs before making the new snapshot visible to readers. + // Persist changed catalogs before making the new snapshot visible to readers. let publish_result: Result<()> = async { if let Some(next) = next { - persist_index_catalogs(&store, region.region_id, &next).await?; + persist_index_catalogs(&store, region.region_id, &next, &stats).await?; publish_index_version(®ion, Arc::new(next)); unpublished.disarm(); } @@ -161,6 +163,7 @@ async fn build_index_version( now_ms: i64, purger: &IndexFilePurger, unpublished: &mut UnpublishedSeriesFiles, + enable_range_index: bool, ) -> Result<(Option, ReconcileStats)> { let mut stats = ReconcileStats::default(); let files = version @@ -207,7 +210,7 @@ async fn build_index_version( if plan.builds.is_empty() && plan.expired_index_ids.is_empty() && stats.removed_range == 0 - && current.range_indexes.len() == visible.len() + && (!enable_range_index || current.range_indexes.len() == visible.len()) { return Ok((None, stats)); } @@ -224,7 +227,8 @@ async fn build_index_version( for (bucket, expected) in plan.builds { // Complete companion indexes independently so a failed series build preserves them. for file in &bucket.files { - if !range_indexes.contains(&file.file_id().file_id()) + if enable_range_index + && !range_indexes.contains(&file.file_id().file_id()) && let Some(file_id) = build_range_index(store, region, version, file.clone()).await? { @@ -239,14 +243,16 @@ async fn build_index_version( series_indexes.insert(expected.index_uuid, series_handle); } // Cover SSTs outside planned aggregate builds, including skipped buckets. - for file in files { - let file_id = file.file_id().file_id(); - if range_indexes.contains(&file_id) { - continue; - } - if let Some(file_id) = build_range_index(store, region, version, file).await? { - stats.built_range += 1; - range_indexes.insert(file_id); + if enable_range_index { + for file in files { + let file_id = file.file_id().file_id(); + if range_indexes.contains(&file_id) { + continue; + } + if let Some(file_id) = build_range_index(store, region, version, file).await? { + stats.built_range += 1; + range_indexes.insert(file_id); + } } } stats.removed_series = current @@ -267,43 +273,48 @@ async fn build_index_version( Ok((Some(next), stats)) } -/// Writes catalogs in a stable order; the two writes are not atomic together. +/// Writes changed catalogs in a stable order; the two writes are not atomic together. async fn persist_index_catalogs( store: &ObjectStore, region_id: RegionId, next: &SeriesIndexVersion, + stats: &ReconcileStats, ) -> Result<()> { - let mut range_entries = next.range_indexes.iter().copied().collect::>(); - range_entries.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); - let mut series_entries = next - .series_indexes - .values() - .map(|handle| handle.entry().clone()) - .collect::>(); - series_entries.sort_unstable_by_key(|entry| { - ( - entry.bucket_start, - entry.bucket_end, - entry.min_file_sequence, - entry.max_file_sequence, + if stats.built_range + stats.removed_range > 0 { + let mut range_entries = next.range_indexes.iter().copied().collect::>(); + range_entries.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); + store_catalog( + store, + &range_catalog_path(region_id), + &RangeIndexCatalog { + indexes: range_entries, + }, ) - }); - store_catalog( - store, - &range_catalog_path(region_id), - &RangeIndexCatalog { - indexes: range_entries, - }, - ) - .await?; - store_catalog( - store, - &series_catalog_path(region_id), - &SeriesIndexCatalog { - indexes: series_entries, - }, - ) - .await?; + .await?; + } + if stats.built_series + stats.removed_series > 0 { + let mut series_entries = next + .series_indexes + .values() + .map(|handle| handle.entry().clone()) + .collect::>(); + series_entries.sort_unstable_by_key(|entry| { + ( + entry.bucket_start, + entry.bucket_end, + entry.min_file_sequence, + entry.max_file_sequence, + ) + }); + store_catalog( + store, + &series_catalog_path(region_id), + &SeriesIndexCatalog { + indexes: series_entries, + }, + ) + .await?; + } Ok(()) } diff --git a/src/mito2/src/series_index/task.rs b/src/mito2/src/series_index/task.rs index a7f263bef33..2afbb245014 100644 --- a/src/mito2/src/series_index/task.rs +++ b/src/mito2/src/series_index/task.rs @@ -77,6 +77,7 @@ pub(crate) fn spawn_series_index_tasks( purge_receiver: UnboundedReceiver, interval: Duration, time_provider: TimeProviderRef, + enable_range_index: bool, ) -> JoinHandle<()> { // Snapshots may retain senders after the worker stops; purge until all senders drop. common_runtime::spawn_compact(run_index_purge_task( @@ -94,6 +95,7 @@ pub(crate) fn spawn_series_index_tasks( state, interval, time_provider, + enable_range_index, } .run() .await; @@ -110,6 +112,7 @@ struct SeriesIndexTask { state: Arc, interval: Duration, time_provider: TimeProviderRef, + enable_range_index: bool, } impl SeriesIndexTask { @@ -154,6 +157,7 @@ impl SeriesIndexTask { self.bucket_width, self.time_provider.current_time_millis(), self.purger.clone(), + self.enable_range_index, ) .await { @@ -211,6 +215,7 @@ mod tests { state: Arc::new(SeriesIndexTaskState::new()), interval: Duration::from_secs(3600), time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, }; task.maintain().await; assert_eq!( diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index cbeb477e04a..c1efe3afdcf 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -154,6 +154,14 @@ impl IndexTest { async fn reconcile( &self, store: ObjectStore, + ) -> crate::error::Result { + self.reconcile_with_range_index(store, true).await + } + + async fn reconcile_with_range_index( + &self, + store: ObjectStore, + enable_range_index: bool, ) -> crate::error::Result { reconcile_series_indexes( 0, @@ -162,11 +170,132 @@ impl IndexTest { Duration::from_secs(100), 0, self.purger.clone(), + enable_range_index, ) .await } } +#[rstest::rstest] +#[case::range_only(vec![1000], 0)] +#[case::series_bucket(vec![1000, 2000, 3000, 4000], 1)] +#[tokio::test] +async fn test_reconcile_range_index_disabled( + #[case] timestamps: Vec, + #[case] expected_series: usize, +) { + use std::sync::Mutex; + + let mut env = TestEnv::with_prefix("series-range-disabled").await; + let (engine, region) = prepare_region_with_timestamps(&mut env, ×tamps).await; + let (test, _receiver) = IndexTest::new(region.clone()); + let writes = Arc::new(Mutex::new(Vec::new())); + let captured = writes.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + captured.lock().unwrap().push(path.to_string()); + inner + })) + .build() + .unwrap(); + let store = test.store.clone().layer(layer); + let range_catalog = range_catalog_path(region.region_id); + let series_catalog = series_catalog_path(region.region_id); + let files = region + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .map(|file| file.meta_ref().clone()) + .collect::>(); + + // Both planned series buckets and SSTs without a series build must skip range builds. + let stats = test + .reconcile_with_range_index(store.clone(), false) + .await + .unwrap(); + assert_eq!( + (0, expected_series), + (stats.built_range, stats.built_series) + ); + assert!(region.series_index_version().range_indexes.is_empty()); + assert!(!store.exists(&range_catalog).await.unwrap()); + assert_eq!( + expected_series > 0, + store.exists(&series_catalog).await.unwrap() + ); + for file in &files { + assert!( + !store + .exists(&range_index_path(region.region_id, file.file_id)) + .await + .unwrap() + ); + } + writes.lock().unwrap().clear(); + let disabled = region.series_index_version(); + test.reconcile_with_range_index(store.clone(), false) + .await + .unwrap(); + assert!(Arc::ptr_eq(&disabled, ®ion.series_index_version())); + assert!(writes.lock().unwrap().is_empty()); + + // Re-enabling fills missing range coverage without rewriting the series catalog. + let stats = test.reconcile(store.clone()).await.unwrap(); + assert_eq!((files.len(), 0), (stats.built_range, stats.built_series)); + assert!(writes.lock().unwrap().contains(&range_catalog)); + assert!(!writes.lock().unwrap().contains(&series_catalog)); + writes.lock().unwrap().clear(); + let enabled = region.series_index_version(); + assert_eq!(files.len(), enabled.range_indexes.len()); + assert_eq!( + disabled + .series_indexes + .keys() + .collect::>(), + enabled.series_indexes.keys().collect() + ); + + // Disabling again retains live entries and does not write either catalog. + test.reconcile_with_range_index(store.clone(), false) + .await + .unwrap(); + assert!(Arc::ptr_eq(&enabled, ®ion.series_index_version())); + assert!(writes.lock().unwrap().is_empty()); + + // Obsolete entries are still pruned and persisted while range builds are disabled. + region.version_control.apply_edit( + Some(crate::manifest::action::RegionEdit { + files_to_remove: files, + files_to_add: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }), + &[], + crate::test_util::new_noop_file_purger(), + ); + let stats = test + .reconcile_with_range_index(store.clone(), false) + .await + .unwrap(); + assert_eq!( + (0, timestamps.len()), + (stats.built_range, stats.removed_range) + ); + assert!(writes.lock().unwrap().contains(&range_catalog)); + assert!(!writes.lock().unwrap().contains(&series_catalog)); + let restored = load_version_control(&store, region.region_id, &test.purger).await; + assert!(restored.current().range_indexes.is_empty()); + writes.lock().unwrap().clear(); + test.reconcile_with_range_index(store, false).await.unwrap(); + assert!(writes.lock().unwrap().is_empty()); + engine.stop().await.unwrap(); +} + #[tokio::test] async fn test_reconcile_restores_and_reuses_indexes() { let mut env = TestEnv::with_prefix("series-reconcile").await; @@ -683,8 +812,11 @@ async fn wait_for(mut ready: impl FnMut() -> bool) { .unwrap(); } +#[rstest::rstest] +#[case::range_enabled(true)] +#[case::range_disabled(false)] #[tokio::test] -async fn test_maintenance_wakeup_and_timer() { +async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { let mut env = TestEnv::with_prefix("series-task").await; let (engine, region) = prepare_region(&mut env).await; let mut options = region.version().options.clone(); @@ -708,8 +840,13 @@ async fn test_maintenance_wakeup_and_timer() { receiver, Duration::from_secs(3600), clock.clone(), + enable_range_index, ); wait_for(|| !region.series_index_version().series_indexes.is_empty()).await; + assert_eq!( + enable_range_index, + !region.series_index_version().range_indexes.is_empty() + ); // Advance event time; periodic maintenance must expire coverage without a wakeup. clock.set_now(201_000); wait_for(|| region.series_index_version().series_indexes.is_empty()).await; @@ -779,6 +916,7 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { receiver, Duration::from_secs(3600), Arc::new(StdTimeProvider), + true, ) .await .unwrap(); diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 31fad8236c9..24eec5dc2b3 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -622,6 +622,7 @@ impl WorkerStarter { purge_receiver, self.config.experimental_series_index_maintenance_interval, self.time_provider.clone(), + self.config.experimental_enable_range_index, ) }); let now = self.time_provider.current_time_millis(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index bd761f39a31..cd614933353 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2510,6 +2510,7 @@ experimental_manifest_keep_removed_file_count = 256 experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_enable_series_index = false +experimental_enable_range_index = false experimental_series_index_maintenance_interval = "5m" experimental_series_index_bucket_width = "5days" experimental_compaction_memory_limit = "unlimited" From 33cfb72a435ffee266acff7d531453458458228c Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Sun, 20 Sep 2026 04:23:37 +0000 Subject: [PATCH 19/63] fix: make database export assertions portable on Windows (#9256) fix: compare database export paths portably on Windows Signed-off-by: jeremyhi --- AGENTS.md | 5 +++ tests-integration/AGENTS.md | 36 +++++++++++++++++++ .../tests/export_logical_tables.rs | 30 ++++++++++------ 3 files changed, 61 insertions(+), 10 deletions(-) create mode 100644 tests-integration/AGENTS.md diff --git a/AGENTS.md b/AGENTS.md index 30e0077caed..ef56c326805 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,6 +66,7 @@ a module map, read/write paths, change-coupling points, and gotchas: - [`src/flow/AGENTS.md`](src/flow/AGENTS.md) - [`src/frontend/AGENTS.md`](src/frontend/AGENTS.md) - [`src/meta-srv/AGENTS.md`](src/meta-srv/AGENTS.md) +- [`tests-integration/AGENTS.md`](tests-integration/AGENTS.md) - [`tests/compatibility/AGENTS.md`](tests/compatibility/AGENTS.md) - [`tests/perf/AGENTS.md`](tests/perf/AGENTS.md) @@ -117,6 +118,10 @@ blast radius requires it. | Query regression harness or DSL | Follow `tests/perf/AGENTS.md` | | Enterprise-gated code | Build/test with `--features enterprise` where applicable and run `make check-enterprise-license` | +For import/export, COPY, or snapshot-storage changes, follow the Windows +portability and validation guidance in +[`tests-integration/AGENTS.md`](tests-integration/AGENTS.md). + ## Before opening a PR 1. If you added a `.rs`, `.py`, or `.ts` file, apply and verify its license diff --git a/tests-integration/AGENTS.md b/tests-integration/AGENTS.md new file mode 100644 index 00000000000..9b83a05497c --- /dev/null +++ b/tests-integration/AGENTS.md @@ -0,0 +1,36 @@ +# tests-integration — Agent & Contributor Guide + +Integration test modules are registered in `tests/main.rs` (`autotests = false`). +Shared instance and storage helpers live in `src/instance.rs` and +`src/test_util.rs`. See [README.md](README.md) for external-service setup. + +## Import/export portability + +Import/export and COPY tests must work on Windows as well as Linux/macOS. +[#9227](https://github.com/GreptimeTeam/greptimedb/issues/9227) records export +tests failing because native Windows paths were compared with `/`-separated +COPY locations. + +- Keep native filesystem paths (`Path`/`PathBuf`), file URIs, and `/`-separated + object-store keys distinct. Compare filesystem paths as paths; for COPY + location assertions, normalize both sides to the defined location format. + Do not compare native `PathBuf::join().display()` strings directly with + `/`-separated export locations. +- Build filesystem fixtures with temporary directories and convert absolute + paths to file URIs with `Url::from_file_path`; do not hard-code `/tmp` or + concatenate `file://` with native paths. Close file handles before rename or + cleanup, including import-state writes. + +## Validation + +Run focused cases from the repository root: + +```bash +cargo nextest run -p tests-integration --test main -E 'test()' +``` + +For import/export changes, run affected tests on Windows. The reference job is +`test-on-windows` in [Nightly CI](../.github/workflows/nightly-ci.yml) +(`cargo nextest run -F dashboard`). Linux/macOS success does not establish +Windows compatibility. If Windows execution is unavailable, explicitly report +that gap and inspect platform-specific paths and file operations. diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 9b4ea97a4a4..8d5ea3e42bd 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::path::PathBuf; use std::sync::Arc; use common_query::OutputData; @@ -292,15 +293,13 @@ async fn database_export_roundtrip(instance: &Arc) { .path() .join("data") .join(format!("{name}.parquet")) - .to_str() - .unwrap() - .to_string() }) .collect::>(); assert_eq!( summary .output_files .into_iter() + .map(PathBuf::from) .collect::>(), expected ); @@ -697,13 +696,24 @@ async fn database_export_preserves_valid_table_names() { } }) .collect::>(); - assert_eq!( - summary - .output_files - .into_iter() - .collect::>(), - expected - ); + let actual = summary + .output_files + .into_iter() + .collect::>(); + if file_url { + assert_eq!(actual, expected); + } else { + assert_eq!( + actual + .into_iter() + .map(PathBuf::from) + .collect::>(), + expected + .into_iter() + .map(PathBuf::from) + .collect::>() + ); + } } for name in &names { let path = directory.join(format!("{name}.parquet")); From 43ce07fa694ca53c87a2f104bfe264e0da2e96e9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 04:58:47 +0000 Subject: [PATCH 20/63] fix(mito2): cancel cache construction for incomplete scans (#9254) * fix(mito2): cancel cache construction for incomplete scans CacheBatchBuffer spawned the background concat task and dropped its join handle. A scan that was cancelled or failed therefore left the task alive: it kept compacting already queued batches, and while waiting for a range result memory permit it held them, even though without a finish command the result can never be put into the cache. Keep the handle and abort it when the buffer is dropped. The handle is cleared once the task owns the finish command, so a completed scan still populates the cache after its stream is dropped. Abort does not preempt a concat that is already running; it takes effect the next time the task is polled. Signed-off-by: Dennis Zhuang * test(mito2): wait for the permit park before cancelling the buffer An empty buffered_batches only proves the batches were enqueued, so the cancellation test could abort a concat task that had never been polled. Count acquisitions that find too few permits, a test-only signal, and drop the buffer once the task has reached that wait. Nothing awaits between the check and the parking, so an observed increment means the caller is about to wait for permits the test holds for the rest of the case. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/cache.rs | 20 ++++++++ src/mito2/src/read/range_cache.rs | 81 ++++++++++++++++++++++++++----- 2 files changed, 89 insertions(+), 12 deletions(-) diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 2e98ef33454..0d43d31f434 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -99,6 +99,10 @@ pub(crate) struct RangeResultMemoryLimiter { semaphore: Arc, permit_bytes: usize, total_permits: usize, + /// Number of acquisitions that found too few permits and parked. Test-only + /// signal to synchronize with a caller waiting inside [Self::acquire]. + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize, } impl Default for RangeResultMemoryLimiter { @@ -120,6 +124,8 @@ impl RangeResultMemoryLimiter { semaphore: Arc::new(tokio::sync::Semaphore::new(total_permits)), permit_bytes, total_permits, + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize::new(0), } } @@ -128,6 +134,12 @@ impl RangeResultMemoryLimiter { self.permit_bytes } + #[cfg(test)] + pub(crate) fn waited_acquires(&self) -> usize { + self.waited_acquires + .load(std::sync::atomic::Ordering::Acquire) + } + #[cfg(test)] pub(crate) fn available_permits(&self) -> usize { self.semaphore.available_permits() @@ -144,6 +156,14 @@ impl RangeResultMemoryLimiter { } .fail(); } + // Nothing awaits between this check and the parking below, so an observed + // increment means the caller is about to wait for the missing permits. + #[cfg(test)] + if self.semaphore.available_permits() < permits { + self.waited_acquires + .fetch_add(1, std::sync::atomic::Ordering::Release); + } + self.semaphore .acquire_many(permits as u32) .await diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index ddf8c3cc356..1108b074697 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -736,26 +736,33 @@ struct CacheBatchBuffer { buffered_rows: usize, buffered_size: usize, sender: Option>, + /// Handle of the background concat task. Cleared once the task owns a finish + /// command, so only unfinished scans cancel it. + concat_task: Option>, } impl CacheBatchBuffer { fn new(cache_strategy: &CacheStrategy) -> Self { - let sender = cache_strategy.range_result_memory_limiter().map(|limiter| { - let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); - let (tx, rx) = mpsc::unbounded_channel(); - common_runtime::spawn_query(run_cache_concat_task( - rx, - limiter.clone(), - skip_threshold_bytes, - )); - tx - }); + let (sender, concat_task) = cache_strategy + .range_result_memory_limiter() + .map(|limiter| { + let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); + let (tx, rx) = mpsc::unbounded_channel(); + let task = common_runtime::spawn_query(run_cache_concat_task( + rx, + limiter.clone(), + skip_threshold_bytes, + )); + (tx, task) + }) + .unzip(); Self { buffered_batches: Vec::new(), buffered_rows: 0, buffered_size: 0, sender, + concat_task, } } @@ -814,9 +821,22 @@ impl CacheBatchBuffer { part_metrics, result_tx, }) - .is_err() + .is_ok() { - self.sender = None; + // The task now owns the finish command, so it may keep populating the + // cache after the scan stream is dropped. + self.concat_task = None; + } + } +} + +impl Drop for CacheBatchBuffer { + fn drop(&mut self) { + // Still holding the handle means no finish command was sent: the scan was + // cancelled or failed and the queued batches can never reach the cache. + // Aborting releases them even while the task waits for a memory permit. + if let Some(task) = &self.concat_task { + task.abort(); } } } @@ -1570,6 +1590,43 @@ mod tests { assert_eq!(replayed[1].num_rows(), 1); } + #[tokio::test] + async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() { + let strategy = test_cache_strategy(); + let limiter = strategy.range_result_memory_limiter().unwrap(); + // Take every permit so the concat task blocks before it can compact. + let _permit = limiter + .acquire(limiter.available_permits() * limiter.permit_bytes()) + .await + .unwrap(); + let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); + let weak = Arc::downgrade(batch.column(0)); + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch.clone()).unwrap(); + buffer.push(batch).unwrap(); + // Both batches were handed to the task instead of staying in the buffer. + assert!(buffer.buffered_batches.is_empty()); + + // Cancel only once the task is parked on the permit, otherwise the abort + // could land on a task that was never polled. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while limiter.waited_acquires() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("concat task must wait for a memory permit"); + drop(buffer); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled cache work must release input without waiting for a permit"); + } + #[tokio::test] async fn cache_batch_buffer_finishes_pending_batches() { let strategy = test_cache_strategy(); From bbc4c4bf1d4a3cbc37074672893eadebf25a54b8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 06:18:24 +0000 Subject: [PATCH 21/63] perf(otlp): share trace resource and scope attributes (#9253) * perf(otlp): share trace resource and scope attributes Parsing an OTLP trace request copied the resource attributes and the scope attributes once per span. A request with wide resource attributes kept one full copy per span alive until the rows were built. Spans and their group now share one `Arc` per resource and per scope. Empty attributes stay `None`, so a resource or scope without attributes costs no allocation and no refcounting. v0 takes ownership back when it encodes, v1 clones attributes item by item instead of rebuilding a Vec, and v2 borrows the span instead of cloning it for every row. Column values, attribute order, the `service.name` lookup and the auxiliary table writes are unchanged. Signed-off-by: Dennis Zhuang * perf(otlp): stop copying v1 attribute keys The v1 row writer rebuilds every column name as `{prefix}.{key}`, so the key string it cloned from the shared resource and scope attributes was dropped unused. `resource_attributes.service.name` was also cloned in full before being skipped for the top level service name column. Shared attributes are now read in place and only values that reach a row are copied. Span attributes keep moving their values as before. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../src/instance/otlp/trace_ingest/tests.rs | 4 +- src/servers/src/otlp/trace/attributes.rs | 34 +++ src/servers/src/otlp/trace/span.rs | 81 ++++-- src/servers/src/otlp/trace/v0.rs | 8 +- src/servers/src/otlp/trace/v1.rs | 272 +++++++++++------- src/servers/src/otlp/trace/v2.rs | 46 +-- 6 files changed, 293 insertions(+), 152 deletions(-) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index c0ba1a91239..1459119efdd 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -340,10 +340,10 @@ fn trace_span(span_id: &str, span_attributes: Vec) -> TraceSpan { trace_id: "trace-id".to_string(), span_id: span_id.to_string(), parent_span_id: None, - resource_attributes: Attributes::from(vec![]), + resource_attributes: Attributes::from(vec![]).into(), scope_name: "scope".to_string(), scope_version: "v1".to_string(), - scope_attributes: Attributes::from(vec![]), + scope_attributes: Attributes::from(vec![]).into(), trace_state: String::new(), span_name: "operation".to_string(), span_kind: "SPAN_KIND_SERVER".to_string(), diff --git a/src/servers/src/otlp/trace/attributes.rs b/src/servers/src/otlp/trace/attributes.rs index 471832b8920..58bc5891540 100644 --- a/src/servers/src/otlp/trace/attributes.rs +++ b/src/servers/src/otlp/trace/attributes.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::fmt::Display; +use std::sync::Arc; use opentelemetry_proto::tonic::common::v1::any_value::Value::{ ArrayValue, BoolValue, BytesValue, DoubleValue, IntValue, KvlistValue, StringValue, @@ -147,6 +148,39 @@ impl Attributes { } } +/// Attributes shared by every span of one resource or scope. +/// +/// Empty attributes are kept as `None` so that the common case of a resource or +/// scope without attributes costs no allocation and no atomic refcounting. +#[derive(Debug, Clone)] +pub struct SharedAttributes(Option>); + +impl From for SharedAttributes { + fn from(attributes: Attributes) -> Self { + Self((!attributes.0.is_empty()).then(|| Arc::new(attributes))) + } +} + +impl AsRef for SharedAttributes { + fn as_ref(&self) -> &Attributes { + static EMPTY: Attributes = Attributes(Vec::new()); + self.0.as_deref().unwrap_or(&EMPTY) + } +} + +impl SharedAttributes { + /// Returns owned attributes, copying them unless this is the last holder. + pub fn into_owned(self) -> Attributes { + match self.0 { + // Spans usually still share with their group, so skip the + // compare-exchange that would fail inside `unwrap_or_clone`. + Some(attributes) if Arc::strong_count(&attributes) > 1 => attributes.as_ref().clone(), + Some(attributes) => Arc::unwrap_or_clone(attributes), + None => Attributes(Vec::new()), + } + } +} + #[cfg(test)] mod tests { use opentelemetry_proto::tonic::common::v1::any_value::Value; diff --git a/src/servers/src/otlp/trace/span.rs b/src/servers/src/otlp/trace/span.rs index 37a6e28935f..a8a51705ffe 100644 --- a/src/servers/src/otlp/trace/span.rs +++ b/src/servers/src/otlp/trace/span.rs @@ -17,13 +17,13 @@ use std::fmt::Display; use common_time::timestamp::Timestamp; use itertools::Itertools; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; -use opentelemetry_proto::tonic::common::v1::{InstrumentationScope, KeyValue, any_value}; +use opentelemetry_proto::tonic::common::v1::{InstrumentationScope, any_value}; use opentelemetry_proto::tonic::trace::v1::span::{Event, Link}; use opentelemetry_proto::tonic::trace::v1::{Span, Status}; use serde::Serialize; use crate::otlp::trace::KEY_SERVICE_NAME; -use crate::otlp::trace::attributes::Attributes; +use crate::otlp::trace::attributes::{Attributes, SharedAttributes}; use crate::otlp::utils::bytes_to_hex_string; #[derive(Debug, Clone)] @@ -35,10 +35,10 @@ pub struct TraceSpan { pub parent_span_id: Option, // the following are fields - pub resource_attributes: Attributes, + pub resource_attributes: SharedAttributes, pub scope_name: String, pub scope_version: String, - pub scope_attributes: Attributes, + pub scope_attributes: SharedAttributes, pub trace_state: String, pub span_name: String, pub span_kind: String, @@ -56,10 +56,10 @@ pub type TraceSpans = Vec; #[derive(Debug, Clone)] pub struct TraceSpanGroup { pub service_name: Option, - pub resource_attributes: Attributes, + pub resource_attributes: SharedAttributes, pub scope_name: String, pub scope_version: String, - pub scope_attributes: Attributes, + pub scope_attributes: SharedAttributes, pub spans: TraceSpans, } @@ -205,8 +205,10 @@ impl SpanEvents { pub fn parse_span( service_name: Option, - resource_attrs: &[KeyValue], - scope: &InstrumentationScope, + resource_attrs: &SharedAttributes, + scope_name: &str, + scope_version: &str, + scope_attrs: &SharedAttributes, span: Span, ) -> TraceSpan { let (span_status_code, span_status_message) = status_to_string(&span.status); @@ -221,12 +223,12 @@ pub fn parse_span( Some(bytes_to_hex_string(&span.parent_span_id)) }, - resource_attributes: Attributes::from(resource_attrs), + resource_attributes: resource_attrs.clone(), trace_state: span.trace_state, - scope_name: scope.name.clone(), - scope_version: scope.version.clone(), - scope_attributes: Attributes::from(scope.attributes.clone()), + scope_name: scope_name.to_string(), + scope_version: scope_version.to_string(), + scope_attributes: scope_attrs.clone(), span_name: span.name, span_kind, @@ -282,23 +284,32 @@ pub fn parse(request: ExportTraceServiceRequest) -> TraceSpanGroups { _ => None, }); + let resource_attrs = SharedAttributes::from(Attributes::from(resource_attrs)); for scope_spans in resource_spans.scope_spans { - let scope = scope_spans.scope.unwrap_or_default(); + let InstrumentationScope { + name: scope_name, + version: scope_version, + attributes: scope_attrs, + .. + } = scope_spans.scope.unwrap_or_default(); + let scope_attrs = SharedAttributes::from(Attributes::from(scope_attrs)); let mut spans = Vec::with_capacity(scope_spans.spans.len()); for span in scope_spans.spans { spans.push(parse_span( service_name.clone(), &resource_attrs, - &scope, + &scope_name, + &scope_version, + &scope_attrs, span, )); } groups.push(TraceSpanGroup { service_name: service_name.clone(), - resource_attributes: Attributes::from(&resource_attrs[..]), - scope_name: scope.name, - scope_version: scope.version, - scope_attributes: Attributes::from(scope.attributes), + resource_attributes: resource_attrs.clone(), + scope_name, + scope_version, + scope_attributes: scope_attrs, spans, }); } @@ -378,6 +389,7 @@ mod tests { ScopeSpans { scope: Some(InstrumentationScope { name: "scope-1".to_string(), + attributes: vec![make_kv("scope.key", "scope-1-value")], ..Default::default() }), spans: vec![make_span(0x11, 0x21), make_span(0x12, 0x22)], @@ -386,6 +398,7 @@ mod tests { ScopeSpans { scope: Some(InstrumentationScope { name: "scope-2".to_string(), + attributes: vec![make_kv("scope.key", "scope-2-value")], ..Default::default() }), spans: vec![make_span(0x13, 0x23)], @@ -421,5 +434,37 @@ mod tests { assert_eq!(groups[1].spans.len(), 1); assert_eq!(groups[2].service_name.as_deref(), Some("svc-b")); assert_eq!(groups[2].scope_name, "scope-3"); + assert_eq!( + groups[0].scope_attributes.as_ref().get_ref(), + &[make_kv("scope.key", "scope-1-value")] + ); + + // Scopes of one resource share a single copy of the resource attributes, + // spans share their group's attributes, and unrelated resources and + // scopes stay separate. + assert!(std::ptr::eq( + groups[0].resource_attributes.as_ref(), + groups[1].resource_attributes.as_ref(), + )); + assert!(!std::ptr::eq( + groups[0].resource_attributes.as_ref(), + groups[2].resource_attributes.as_ref(), + )); + assert!(!std::ptr::eq( + groups[0].scope_attributes.as_ref(), + groups[1].scope_attributes.as_ref(), + )); + for group in &groups { + for span in &group.spans { + assert!(std::ptr::eq( + group.resource_attributes.as_ref(), + span.resource_attributes.as_ref() + )); + assert!(std::ptr::eq( + group.scope_attributes.as_ref(), + span.scope_attributes.as_ref() + )); + } + } } } diff --git a/src/servers/src/otlp/trace/v0.rs b/src/servers/src/otlp/trace/v0.rs index f3f7bdc29fa..058052a1857 100644 --- a/src/servers/src/otlp/trace/v0.rs +++ b/src/servers/src/otlp/trace/v0.rs @@ -152,14 +152,14 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> row_writer::write_json( writer, "scope_attributes", - span.scope_attributes.into(), + span.scope_attributes.into_owned().into(), &mut row, )?; row_writer::write_json( writer, "resource_attributes", - span.resource_attributes.into(), + span.resource_attributes.into_owned().into(), &mut row, )?; @@ -243,10 +243,10 @@ mod tests { trace_id: trace_id.to_string(), span_id: span_id.to_string(), parent_span_id: None, - resource_attributes: Attributes::from(vec![]), + resource_attributes: Attributes::from(vec![]).into(), scope_name: "scope".to_string(), scope_version: "v1".to_string(), - scope_attributes: Attributes::from(vec![]), + scope_attributes: Attributes::from(vec![]).into(), trace_state: String::new(), span_name: "op".to_string(), span_kind: "SPAN_KIND_SERVER".to_string(), diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index bd556c4f056..615ae472183 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -23,6 +23,7 @@ use api::v1::{ }; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_grpc::precision::Precision; +use opentelemetry_proto::tonic::common::v1::KeyValue; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; use crate::error::Result; @@ -436,23 +437,23 @@ fn write_span_to_row_inner( write_attributes_with_schema( writer, "span_attributes", - span.span_attributes, + span.span_attributes.take(), &mut row, row_index, batch_schema.as_deref_mut(), )?; - write_attributes_with_schema( + write_shared_attributes_with_schema( writer, "scope_attributes", - span.scope_attributes, + span.scope_attributes.as_ref(), &mut row, row_index, batch_schema.as_deref_mut(), )?; - write_attributes_with_schema( + write_shared_attributes_with_schema( writer, "resource_attributes", - span.resource_attributes, + span.resource_attributes.as_ref(), &mut row, row_index, batch_schema, @@ -535,124 +536,177 @@ pub(crate) fn write_attributes( row: &mut Vec, ) -> Result<()> { let row_index = writer.num_rows(); - write_attributes_with_schema(writer, prefix, attributes, row, row_index, None) + write_attributes_with_schema(writer, prefix, attributes.take(), row, row_index, None) } -/// Writes flattened attributes without coercion and optionally records their actual types. +/// Skips `resource_attributes.service.name` because it is already copied to the +/// top level as `SERVICE_NAME_COLUMN`. +fn skipped_attribute(prefix: &str, key: &str) -> bool { + prefix == "resource_attributes" && key == KEY_SERVICE_NAME +} + +/// Writes flattened attributes owned by one span. fn write_attributes_with_schema( writer: &mut TableData, prefix: &str, - attributes: Attributes, + attributes: Vec, row: &mut Vec, row_index: usize, mut batch_schema: Option<&mut TraceBatchSchema>, ) -> Result<()> { - for attr in attributes.take().into_iter() { - let key_suffix = attr.key; - // skip resource_attributes.service.name because its already copied to - // top level as `SERVICE_NAME_COLUMN` - if prefix == "resource_attributes" && key_suffix == KEY_SERVICE_NAME { + for KeyValue { key, value, .. } in attributes { + if skipped_attribute(prefix, &key) { continue; } - - let key = format!("{}.{}", prefix, key_suffix); - match attr.value.and_then(|v| v.value) { - Some(OtlpValue::StringValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_value_type(&key, row_index, ColumnDataType::String); - } - // Keep the raw request value here. Mixed trace types are reconciled later - // in the frontend once we can also see the existing table schema. - writer.write_field_unchecked( - &key, - ColumnDataType::String, - Some(ValueData::StringValue(v)), - row, - ); - } - Some(OtlpValue::BoolValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_value_type(&key, row_index, ColumnDataType::Boolean); - } - // Do not coerce or promote types while building the request-local rows. - writer.write_field_unchecked( - &key, - ColumnDataType::Boolean, - Some(ValueData::BoolValue(v)), - row, - ); - } - Some(OtlpValue::IntValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_value_type(&key, row_index, ColumnDataType::Int64); - } - // Preserving the original value avoids order-dependent behavior inside one batch. - writer.write_field_unchecked( - &key, - ColumnDataType::Int64, - Some(ValueData::I64Value(v)), - row, - ); - } - Some(OtlpValue::DoubleValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_value_type(&key, row_index, ColumnDataType::Float64); - } - writer.write_field_unchecked( - &key, - ColumnDataType::Float64, - Some(ValueData::F64Value(v)), - row, - ); - } - Some(OtlpValue::ArrayValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); - } - writer.write_column_unchecked( - row_writer::build_json_column_schema(key), - Some(ValueData::BinaryValue( - any_value_to_jsonb(OtlpValue::ArrayValue(v)).to_vec(), - )), - row, - ); - } - Some(OtlpValue::KvlistValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); - } - writer.write_column_unchecked( - row_writer::build_json_column_schema(key), - Some(ValueData::BinaryValue( - any_value_to_jsonb(OtlpValue::KvlistValue(v)).to_vec(), - )), - row, - ); - } - Some(OtlpValue::BytesValue(v)) => { - if let Some(batch_schema) = batch_schema.as_deref_mut() { - batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Binary); - } - writer.write_field_unchecked( - key, - ColumnDataType::Binary, - Some(ValueData::BinaryValue(v)), - row, - ); - } - // `StringValueStrindex` is profiling-signal-only and references the - // Profiling `ProfilesDictionary.string_table`, which is unavailable to - // traces. Per the OTLP spec, non-Profiling receivers must treat it as a - // non-fatal issue and process the value as if it were absent. Like the - // `None` arm, no field is written for the attribute. - Some(OtlpValue::StringValueStrindex(_)) => {} - None => {} - } + write_attribute_with_schema( + writer, + prefix, + &key, + value.and_then(|v| v.value), + row, + row_index, + batch_schema.as_deref_mut(), + ); } Ok(()) } +/// Writes flattened attributes shared by every span of a resource or scope. +/// +/// The column name is rebuilt from `prefix`, so keys are read in place and only +/// values that reach a row are copied. +fn write_shared_attributes_with_schema( + writer: &mut TableData, + prefix: &str, + attributes: &Attributes, + row: &mut Vec, + row_index: usize, + mut batch_schema: Option<&mut TraceBatchSchema>, +) -> Result<()> { + for attr in attributes.get_ref() { + if skipped_attribute(prefix, &attr.key) { + continue; + } + write_attribute_with_schema( + writer, + prefix, + &attr.key, + attr.value.as_ref().and_then(|v| v.value.clone()), + row, + row_index, + batch_schema.as_deref_mut(), + ); + } + + Ok(()) +} + +/// Writes one flattened attribute without coercion and optionally records its actual type. +fn write_attribute_with_schema( + writer: &mut TableData, + prefix: &str, + key_suffix: &str, + value: Option, + row: &mut Vec, + row_index: usize, + batch_schema: Option<&mut TraceBatchSchema>, +) { + let key = format!("{}.{}", prefix, key_suffix); + match value { + Some(OtlpValue::StringValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::String); + } + // Keep the raw request value here. Mixed trace types are reconciled later + // in the frontend once we can also see the existing table schema. + writer.write_field_unchecked( + &key, + ColumnDataType::String, + Some(ValueData::StringValue(v)), + row, + ); + } + Some(OtlpValue::BoolValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Boolean); + } + // Do not coerce or promote types while building the request-local rows. + writer.write_field_unchecked( + &key, + ColumnDataType::Boolean, + Some(ValueData::BoolValue(v)), + row, + ); + } + Some(OtlpValue::IntValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Int64); + } + // Preserving the original value avoids order-dependent behavior inside one batch. + writer.write_field_unchecked( + &key, + ColumnDataType::Int64, + Some(ValueData::I64Value(v)), + row, + ); + } + Some(OtlpValue::DoubleValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Float64); + } + writer.write_field_unchecked( + &key, + ColumnDataType::Float64, + Some(ValueData::F64Value(v)), + row, + ); + } + Some(OtlpValue::ArrayValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::ArrayValue(v)).to_vec(), + )), + row, + ); + } + Some(OtlpValue::KvlistValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::KvlistValue(v)).to_vec(), + )), + row, + ); + } + Some(OtlpValue::BytesValue(v)) => { + if let Some(batch_schema) = batch_schema { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Binary); + } + writer.write_field_unchecked( + key, + ColumnDataType::Binary, + Some(ValueData::BinaryValue(v)), + row, + ); + } + // `StringValueStrindex` is profiling-signal-only and references the + // Profiling `ProfilesDictionary.string_table`, which is unavailable to + // traces. Per the OTLP spec, non-Profiling receivers must treat it as a + // non-fatal issue and process the value as if it were absent. Like the + // `None` arm, no field is written for the attribute. + Some(OtlpValue::StringValueStrindex(_)) => {} + None => {} + } +} + #[cfg(test)] mod tests { use api::v1::value::ValueData; @@ -679,10 +733,10 @@ mod tests { trace_id: trace_id.to_string(), span_id: span_id.to_string(), parent_span_id: None, - resource_attributes: Attributes::from(vec![]), + resource_attributes: Attributes::from(vec![]).into(), scope_name: "scope".to_string(), scope_version: "v1".to_string(), - scope_attributes: Attributes::from(vec![]), + scope_attributes: Attributes::from(vec![]).into(), trace_state: String::new(), span_name: "op".to_string(), span_kind: "SPAN_KIND_SERVER".to_string(), diff --git a/src/servers/src/otlp/trace/v2.rs b/src/servers/src/otlp/trace/v2.rs index 6d4ee86e326..d4628c3b884 100644 --- a/src/servers/src/otlp/trace/v2.rs +++ b/src/servers/src/otlp/trace/v2.rs @@ -48,13 +48,13 @@ pub(super) fn v2_to_grpc_main_insert_requests( /// Builds the fixed row-oriented payload for the main v2 trace table. fn build_trace_table_data(spans: &[TraceSpan]) -> Result { let mut writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); - for span in spans.iter().cloned() { + for span in spans { write_span_to_row(&mut writer, span)?; } Ok(writer) } -fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { +fn write_span_to_row(writer: &mut TableData, span: &TraceSpan) -> Result<()> { ensure!( span.start_in_nanosecond <= i64::MAX as u64, TimestampOverflowSnafu { @@ -89,18 +89,21 @@ fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { make_column_data( DURATION_NANO_COLUMN, ColumnDataType::Int64, - Some(ValueData::I64Value(span_duration_nano(&span))), + Some(ValueData::I64Value(span_duration_nano(span))), ), - make_string_column_data(PARENT_SPAN_ID_COLUMN, span.parent_span_id), - make_string_column_data(TRACE_ID_COLUMN, Some(span.trace_id)), - make_string_column_data(SPAN_ID_COLUMN, Some(span.span_id)), - make_string_column_data(SPAN_KIND_COLUMN, Some(span.span_kind)), - make_string_column_data(SPAN_NAME_COLUMN, Some(span.span_name)), - make_string_column_data(SPAN_STATUS_CODE, Some(span.span_status_code)), - make_string_column_data(SPAN_STATUS_MESSAGE_COLUMN, Some(span.span_status_message)), - make_string_column_data(TRACE_STATE_COLUMN, Some(span.trace_state)), - make_string_column_data(SCOPE_NAME_COLUMN, Some(span.scope_name)), - make_string_column_data(SCOPE_VERSION_COLUMN, Some(span.scope_version)), + make_string_column_data(PARENT_SPAN_ID_COLUMN, span.parent_span_id.clone()), + make_string_column_data(TRACE_ID_COLUMN, Some(span.trace_id.clone())), + make_string_column_data(SPAN_ID_COLUMN, Some(span.span_id.clone())), + make_string_column_data(SPAN_KIND_COLUMN, Some(span.span_kind.clone())), + make_string_column_data(SPAN_NAME_COLUMN, Some(span.span_name.clone())), + make_string_column_data(SPAN_STATUS_CODE, Some(span.span_status_code.clone())), + make_string_column_data( + SPAN_STATUS_MESSAGE_COLUMN, + Some(span.span_status_message.clone()), + ), + make_string_column_data(TRACE_STATE_COLUMN, Some(span.trace_state.clone())), + make_string_column_data(SCOPE_NAME_COLUMN, Some(span.scope_name.clone())), + make_string_column_data(SCOPE_VERSION_COLUMN, Some(span.scope_version.clone())), ] .into_iter(), &mut row, @@ -114,7 +117,7 @@ fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { semantic_type: SemanticType::Tag as i32, ..Default::default() }, - span.service_name.map(ValueData::StringValue), + span.service_name.clone().map(ValueData::StringValue), )), &mut row, )?; @@ -122,29 +125,34 @@ fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { row_writer::write_json2( writer, SPAN_ATTRIBUTES_COLUMN, - span.span_attributes, + &span.span_attributes, &mut row, )?; row_writer::write_json2( writer, SCOPE_ATTRIBUTES_COLUMN, - span.scope_attributes, + span.scope_attributes.as_ref(), &mut row, )?; row_writer::write_json2( writer, RESOURCE_ATTRIBUTES_COLUMN, - span.resource_attributes, + span.resource_attributes.as_ref(), &mut row, )?; row_writer::write_json( writer, SPAN_EVENTS_COLUMN, - span.span_events.into(), + span.span_events.clone().into(), + &mut row, + )?; + row_writer::write_json( + writer, + "span_links", + span.span_links.clone().into(), &mut row, )?; - row_writer::write_json(writer, "span_links", span.span_links.into(), &mut row)?; writer.add_row(row); Ok(()) From e82cb0af748132fdfb0c22909b9c92249b83468d Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Sun, 20 Sep 2026 07:31:04 +0000 Subject: [PATCH 22/63] feat(json2): respect type hints during JSON2 type concretization (#9222) * fix: prefer JSON2 type hints for uncast read pushdown * feat(query): materialize JSON get result types before planning * fix(query): apply JSON2 type hints to parsed paths * fix(query): apply JSON2 type hints before distributed planning * chore: remove f * refactor: the code style of json_get_type_hint * chore: add sqlness case * fix: add json expr planner, remove json get type hint * chore: add sqlness test * fix(query): respect JSON2 type hints in query planning * chore: add more sqlness cases * fix: cargo clippy * fix: cargo clippy * chore: update sqlness test result --- src/query/src/datafusion/json_expr_planner.rs | 87 +++++++++++-- .../src/optimizer/json_type_concretize.rs | 54 +++++--- src/query/src/query_engine/state.rs | 1 - .../standalone/common/types/json/json2.result | 32 ++--- .../json/json2_query_respect_type_hint.result | 122 ++++++++++++++++++ .../json/json2_query_respect_type_hint.sql | 55 ++++++++ .../common/types/json/json2_type_hints.result | 14 +- 7 files changed, 312 insertions(+), 53 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result create mode 100644 tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index c934565affa..4ec1b729007 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -15,10 +15,13 @@ use std::sync::{Arc, LazyLock}; use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; use common_function::scalars::json::json_get::JsonGetWithType; use common_function::scalars::udf::create_udf; use datafusion_common::arrow::datatypes::DataType; -use datafusion_common::{Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference}; +use datafusion_common::{ + Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference, plan_datafusion_err, +}; use datafusion_expr::expr::{BinaryExpr, ScalarFunction}; use datafusion_expr::planner::{ ExprPlanner, PlannerResult, RawAggregateExpr, RawBinaryExpr, RawFieldAccessExpr, RawScalarExpr, @@ -28,7 +31,10 @@ use datafusion_expr::type_coercion::functions::{UDFCoercionExt, fields_with_udf} use datafusion_expr::{ Expr, ExprSchemable, GetFieldAccess, Operator, ScalarUDF, WindowFunctionDefinition, }; -use datatypes::extension::json::is_json2_extension_type; +use datatypes::extension::json::{ + Json2ExtensionType, is_json2_extension_type, parse_legacy_json2_settings, +}; +use datatypes::types::json_type::JsonNativeType; use sqlparser::ast::BinaryOperator; /// Rewrites JSON-aware SQL expressions into DataFusion expressions. @@ -142,14 +148,19 @@ impl ExprPlanner for JsonExprPlanner { path.push_str(&json_path_field(name)?); } + let mut args = vec![ + Expr::Column(Column::from((qualifier, field))), + Expr::Literal(ScalarValue::Utf8(Some(path)), None), + ]; + if let Some(json_type) = json_type_hint(field, nested_names)? { + args.push(Expr::Literal( + ScalarValue::try_new_null(&json_type.as_arrow_type())?, + None, + )); + } + Ok(PlannerResult::Planned(Expr::ScalarFunction( - ScalarFunction::new_udf( - json_get, - vec![ - Expr::Column(Column::from((qualifier, field))), - Expr::Literal(ScalarValue::Utf8(Some(path)), None), - ], - ), + ScalarFunction::new_udf(json_get, args), ))) } @@ -185,6 +196,27 @@ impl ExprPlanner for JsonExprPlanner { } } +/// Returns the configured native type for an exact JSON2 object path. +fn json_type_hint(field: &Field, path: &[String]) -> Result> { + let settings = if field.extension_type_name() == Some(Json2ExtensionType::NAME) { + let extension = field + .try_extension_type::() + .map_err(|e| plan_datafusion_err!("invalid JSON2 extension metadata: {e}"))?; + Some(extension.metadata().json_settings().clone()) + } else { + parse_legacy_json2_settings(field.metadata()) + .map_err(|e| plan_datafusion_err!("invalid JSON2 extension metadata: {e}"))? + }; + + Ok(settings.and_then(|settings| { + settings + .type_hints() + .iter() + .find(|hint| hint.path == path) + .map(|hint| JsonNativeType::from(&hint.data_type)) + })) +} + /// Quotes field names containing JSONPath punctuation, preserving literal keys. fn json_path_field(name: &str) -> Result { if !name.is_empty() && name.chars().all(|c| c.is_alphanumeric() || c == '_') { @@ -423,7 +455,9 @@ mod tests { use datafusion_expr::WindowFrame; use datafusion_functions::core::coalesce; use datafusion_functions::math::{abs, power}; - use datatypes::extension::json::Json2ExtensionType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::{JsonSettings, JsonTypeHint}; + use datatypes::prelude::ConcreteDataType; use super::*; @@ -638,6 +672,39 @@ mod tests { Ok(()) } + #[test] + fn test_plan_compound_identifier_applies_json2_type_hint() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["payload".to_string(), "cpu".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + None, + ) + .unwrap(); + let field = Field::new("labels", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))); + + let PlannerResult::Planned(Expr::ScalarFunction(func)) = JsonExprPlanner + .plan_compound_identifier( + &field, + Some(&TableReference::bare("events")), + &["payload".to_string(), "cpu".to_string()], + )? + else { + unreachable!() + }; + + assert_eq!( + Some(DataType::Int64), + extract_json_get_type(&Expr::ScalarFunction(func)) + ); + Ok(()) + } + #[test] fn test_plan_functions() -> Result<()> { let planner = JsonExprPlanner; diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 9051ec0a05c..479dd2395b8 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -15,7 +15,6 @@ use std::any::Any; use std::collections::HashMap; -use arrow_schema::DataType; use common_function::scalars::json::json_get::{JsonGetWithType, parse_json_get_path}; use datafusion::datasource::{DefaultTableSource, TableProvider}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -30,8 +29,10 @@ use table::table::adapter::DfTableProviderAdapter; use crate::dummy_catalog::DummyTableProvider; /// Concretize (deduce) the expected JSON type from query. -/// For example, we can concretize a JSON type of `{ a: { b: Number } }` from `select j.a.b::Int64`. -/// The JSON type will be later set into the scan request, for converting the JSON arrays. +/// +/// For example, we can concretize a JSON type of `{ a: { b: Number } }` from +/// `select j.a.b::Int64`. The JSON type will be later set into the scan request, +/// for converting the JSON arrays. #[derive(Debug)] pub(crate) struct JsonTypeConcretizeRule; @@ -108,7 +109,6 @@ fn apply_json_type_hint( pub(crate) fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); - // JSON2 columns in the final output must retain their complete values even when // predicates or other expressions access only specific paths. // For example, `SELECT j FROM t WHERE json_get(j, 'a') = 1`. @@ -168,13 +168,7 @@ fn deduce_json_type(expr: &Expr) -> Result> { ); }; - let Some(path) = f - .args - .get(1) - .and_then(|expr| expr.as_literal()) - .and_then(|x| x.try_as_str()) - .flatten() - else { + let Some(path) = json_get_path(f) else { return plan_err!( "Second argument of {} is expected to be a string literal, actual: {:?}", JsonGetWithType::NAME, @@ -182,10 +176,10 @@ fn deduce_json_type(expr: &Expr) -> Result> { ); }; - let path = parse_json_get_path(path) + let json_path = parse_json_get_path(path) .map_err(|e| plan_datafusion_err!("Invalid JSONPath {path:?}: {e}"))?; - if path + if json_path .paths .iter() .all(|segment| matches!(segment, Path::Root)) @@ -198,12 +192,14 @@ fn deduce_json_type(expr: &Expr) -> Result> { .get(2) .and_then(|expr| expr.as_literal()) .map(|x| x.data_type()) - .unwrap_or(DataType::Utf8View); - let with_type = - JsonNativeType::try_from(&with_type).map_err(|e| plan_datafusion_err!("{e:?}"))?; + .map(|with_type| { + JsonNativeType::try_from(&with_type).map_err(|e| plan_datafusion_err!("{e:?}")) + }) + .transpose()? + .unwrap_or(JsonNativeType::String); let mut root = with_type; - for segment in path.paths.into_iter().rev() { + for segment in json_path.paths.into_iter().rev() { let name = match segment { Path::Root => continue, Path::DotField(name) | Path::ColonField(name) | Path::ObjectField(name) => name, @@ -219,11 +215,22 @@ fn deduce_json_type(expr: &Expr) -> Result> { Ok(Some((column.name.clone(), root))) } +/// Returns the literal JSON path argument of a `json_get` call. +fn json_get_path(function: &datafusion_expr::expr::ScalarFunction) -> Option<&str> { + function + .args + .get(1) + .and_then(|expr| expr.as_literal()) + .and_then(|value| value.try_as_str()) + .flatten() +} + #[cfg(test)] mod tests { use std::sync::Arc; use api::v1::SemanticType; + use arrow_schema::DataType; use common_function::scalars::udf::create_udf; use datafusion::datasource::provider_as_source; use datafusion::functions_aggregate::expr_fn::count; @@ -231,7 +238,8 @@ mod tests { use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{LogicalPlanBuilder, col, lit}; use datafusion_optimizer::OptimizerContext; - use datatypes::extension::json::Json2ExtensionType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::schema::ColumnSchema; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::{ConcreteDataType, RegionId}; @@ -264,6 +272,12 @@ mod tests { } fn build_json2_scan() -> Result<(Arc, LogicalPlanBuilder)> { + build_json2_scan_with_settings(JsonSettings::default()) + } + + fn build_json2_scan_with_settings( + settings: JsonSettings, + ) -> Result<(Arc, LogicalPlanBuilder)> { let region_id = RegionId::new(1024, 2); let mut builder = RegionMetadataBuilder::new(region_id); let mut json_column = ColumnSchema::new( @@ -271,7 +285,9 @@ mod tests { ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), true, ); - json_column.with_extension_type(&Json2ExtensionType::default()); + json_column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))); builder .push_column_metadata(ColumnMetadata { column_schema: json_column, diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 68a301adf7d..aa3b9417b14 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -208,7 +208,6 @@ impl QueryEngineState { FUNCTION_REGISTRY.function_rewrites(), )), ); - if with_dist_planner { analyzer.rules.push(Arc::new(DistPlannerAnalyzer)); analyzer.rules.push(Arc::new(JsonSchemaConcretizeRule)); diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index ab529a20971..ce0c9748ce1 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -473,14 +473,14 @@ select from json2_finite_paths order by ts; -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ -| ts | j | json_get(json2_finite_paths.j,Utf8("$.hint")) | alpha | beta | json_get(json2_finite_paths.j,Utf8("$.conflict")) | -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ -| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | -| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | -| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | -| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ +| ts | j | json_get(json2_finite_paths.j,Utf8("$.hint"),Utf8View(NULL)) | alpha | beta | json_get(json2_finite_paths.j,Utf8("$.conflict")) | ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | +| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | +| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | +| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ admin compact_table('json2_finite_paths'); @@ -500,14 +500,14 @@ select from json2_finite_paths order by ts; -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ -| ts | j | json_get(json2_finite_paths.j,Utf8("$.hint")) | alpha | beta | json_get(json2_finite_paths.j,Utf8("$.conflict")) | -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ -| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | -| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | -| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | -| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | -+-------------------------+-------------------------------------------+-----------------------------------------------+-------+------+---------------------------------------------------+ ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ +| ts | j | json_get(json2_finite_paths.j,Utf8("$.hint"),Utf8View(NULL)) | alpha | beta | json_get(json2_finite_paths.j,Utf8("$.conflict")) | ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | +| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | +| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | +| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | ++-------------------------+-------------------------------------------+--------------------------------------------------------------+-------+------+---------------------------------------------------+ drop table json2_finite_paths; diff --git a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result new file mode 100644 index 00000000000..dcbffbf6058 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.result @@ -0,0 +1,122 @@ +CREATE TABLE json2_query_respect_type_hint ( + ts TIMESTAMP TIME INDEX, + j JSON2 ( + a BIGINT, + b STRING + ), + exponent DOUBLE +) WITH ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +INSERT INTO json2_query_respect_type_hint VALUES + (1, '{"a":9007199254740993,"b":"three"}', 2), + (2, '{"a":7,"b":"7"}', 2); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_query_respect_type_hint'); + ++----------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_query_respect_type_hint') | ++----------------------------------------------------+ +| 0 | ++----------------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a AS hinted_bigint +FROM json2_query_respect_type_hint; + ++---------------+--------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS hinted_bigint | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+--------------------------------------------------------------------------------------------------+ + +-- Uncast paths use their JSON2 type hints. +SELECT + j.a AS hinted_bigint, + j.b AS hinted_string +FROM json2_query_respect_type_hint +ORDER BY ts; + ++------------------+---------------+ +| hinted_bigint | hinted_string | ++------------------+---------------+ +| 9007199254740993 | three | +| 7 | 7 | ++------------------+---------------+ + +-- A JSON2 type hint takes precedence over an explicit SQL cast for the native +-- read type. +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a::STRING +FROM json2_query_respect_type_hint; + ++---------------+-----------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-----------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.a"), Int64(NULL)) AS Utf8) | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+-----------------------------------------------------------------------------------------------+ + +-- An explicit CAST changes the expression result type. +SELECT + j.a::STRING AS cast_string +FROM json2_query_respect_type_hint +ORDER BY ts; + ++------------------+ +| cast_string | ++------------------+ +| 9007199254740993 | +| 7 | ++------------------+ + +-- A JSON2 type hint takes precedence over expression-context inference. +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT coalesce(j.b, exponent) +FROM json2_query_respect_type_hint; + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: CASE WHEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) IS NOT NULL THEN CAST(json_get(json2_query_respect_type_hint.j, Utf8("$.b"), Utf8View(NULL)) AS Float64) ELSE json2_query_respect_type_hint.exponent END AS coalesce(json_get(json2_query_respect_type_hint.j,Utf8("$.b"),Utf8View(NULL)),json2_query_respect_type_hint.exponent) | +| | TableScan: json2_query_respect_type_hint | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- A STRING type hint is respected, so numeric arithmetic requires an explicit cast. +SELECT j.b + exponent +FROM json2_query_respect_type_hint; + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Cannot coerce arithmetic expression Utf8View + Float64 to valid types + +-- An explicit cast changes the expression type while preserving the STRING native read type. +-- The invalid string value still reports a normal SQL cast error. +SELECT j.b::DOUBLE + exponent +FROM json2_query_respect_type_hint; + +Error: 3001(EngineExecuteQuery), Cast error: Cannot cast string 'three' to value of Float64 type + +DROP TABLE json2_query_respect_type_hint; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql new file mode 100644 index 00000000000..604f7d1f31b --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_query_respect_type_hint.sql @@ -0,0 +1,55 @@ +CREATE TABLE json2_query_respect_type_hint ( + ts TIMESTAMP TIME INDEX, + j JSON2 ( + a BIGINT, + b STRING + ), + exponent DOUBLE +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO json2_query_respect_type_hint VALUES + (1, '{"a":9007199254740993,"b":"three"}', 2), + (2, '{"a":7,"b":"7"}', 2); + +ADMIN FLUSH_TABLE('json2_query_respect_type_hint'); + +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a AS hinted_bigint +FROM json2_query_respect_type_hint; + +-- Uncast paths use their JSON2 type hints. +SELECT + j.a AS hinted_bigint, + j.b AS hinted_string +FROM json2_query_respect_type_hint +ORDER BY ts; + +-- A JSON2 type hint takes precedence over an explicit SQL cast for the native +-- read type. +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT j.a::STRING +FROM json2_query_respect_type_hint; + +-- An explicit CAST changes the expression result type. +SELECT + j.a::STRING AS cast_string +FROM json2_query_respect_type_hint +ORDER BY ts; + +-- A JSON2 type hint takes precedence over expression-context inference. +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT coalesce(j.b, exponent) +FROM json2_query_respect_type_hint; + +-- A STRING type hint is respected, so numeric arithmetic requires an explicit cast. +SELECT j.b + exponent +FROM json2_query_respect_type_hint; + +-- An explicit cast changes the expression type while preserving the STRING native read type. +-- The invalid string value still reports a normal SQL cast error. +SELECT j.b::DOUBLE + exponent +FROM json2_query_respect_type_hint; + +DROP TABLE json2_query_respect_type_hint; diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.result b/tests/cases/standalone/common/types/json/json2_type_hints.result index 7739d1d1f8a..1a892b6e4cb 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.result +++ b/tests/cases/standalone/common/types/json/json2_type_hints.result @@ -51,13 +51,13 @@ SELECT FROM json2_type_hints ORDER BY ts; -+-------------------------------------------------+--------------------------------------------------+----------------------------------------------------+----------------------------------------------+ -| json_get(json2_type_hints.j,Utf8("$.user.age")) | json_get(json2_type_hints.j,Utf8("$.user.name")) | json_get(json2_type_hints.j,Utf8("$.user.active")) | json_get(json2_type_hints.j,Utf8("$.score")) | -+-------------------------------------------------+--------------------------------------------------+----------------------------------------------------+----------------------------------------------+ -| 42 | Alice | true | 3.25 | -| | Bob | | | -| | | | | -+-------------------------------------------------+--------------------------------------------------+----------------------------------------------------+----------------------------------------------+ ++-------------------------------------------------------------+-----------------------------------------------------------------+------------------------------------------------------------------+------------------------------------------------------------+ +| json_get(json2_type_hints.j,Utf8("$.user.age"),Int64(NULL)) | json_get(json2_type_hints.j,Utf8("$.user.name"),Utf8View(NULL)) | json_get(json2_type_hints.j,Utf8("$.user.active"),Boolean(NULL)) | json_get(json2_type_hints.j,Utf8("$.score"),Float64(NULL)) | ++-------------------------------------------------------------+-----------------------------------------------------------------+------------------------------------------------------------------+------------------------------------------------------------+ +| 42 | Alice | true | 3.25 | +| | Bob | | | +| | | | | ++-------------------------------------------------------------+-----------------------------------------------------------------+------------------------------------------------------------------+------------------------------------------------------------+ INSERT INTO json2_type_hints VALUES (4, '{"user":{"age":"bad"}}'); From f40d20110c8525190403f50783a920f81d5a02b8 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Sun, 20 Sep 2026 07:45:37 +0000 Subject: [PATCH 23/63] feat(mito2): split SWCS output files by size threshold (#9259) * feat(mito2): split SWCS output files by size threshold Signed-off-by: Lei, HUANG * fix(mito2): use resolved SWCS output size options Read the SWCS output file size threshold from the resolved region options so database-level compaction settings are honored. Align the existing picker test with this configuration source. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/compaction/window.rs | 45 ++++++++-- src/mito2/src/engine/compaction_test.rs | 115 ++++++++++++++++++++++++ src/mito2/src/region/options.rs | 4 +- 3 files changed, 157 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index e6178054ad2..b719543f9af 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -29,12 +29,12 @@ use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; use crate::compaction::last_non_null::inputs_precede_memtables; use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::error::{JoinSnafu, Result}; -use crate::region::options::MergeMode; +use crate::region::options::{CompactionOptions, MergeMode}; use crate::sst::file::FileHandle; -/// Compaction picker that splits the time range of all involved files to windows, and merges -/// the data segments intersects with those windows of files together so that the output files -/// never overlaps. +/// Compaction picker that splits input files into disjoint time windows and merges the +/// segments within each window. Outputs may be split further by the configured file size +/// threshold, at series boundaries for regions with a primary key. #[derive(Clone, Debug)] pub struct WindowedCompactionPicker { compaction_time_window_seconds: Option, @@ -145,6 +145,13 @@ impl Picker for WindowedCompactionPicker { let picker = self.clone(); let region_id = compaction_region.current_version.metadata.region_id; let current_version = compaction_region.current_version.clone(); + let CompactionOptions::Twcs(options) = &compaction_region.region_options.compaction; + // Use the same output size threshold as TWCS, including zero meaning unlimited. + let max_file_size = options + .max_output_file_size + .map(|size| size.as_bytes()) + .filter(|size| *size > 0) + .map(|size| size as usize); let (outputs, expired_ssts, time_window) = common_runtime::spawn_blocking_compact(move || { picker.pick_inner(region_id, ¤t_version, Timestamp::current_millis()) @@ -156,7 +163,7 @@ impl Picker for WindowedCompactionPicker { outputs, expired_ssts, time_window_size: time_window, - max_file_size: None, // todo (hl): we may need to support `max_file_size` parameter in manual compaction. + max_file_size, })) } } @@ -303,13 +310,16 @@ mod tests { use std::sync::Arc; use std::time::Duration; + use common_base::readable_size::ReadableSize; use common_time::Timestamp; use common_time::range::TimestampRange; use store_api::storage::{FileId, RegionId}; use crate::compaction::compactor::CompactionVersion; + use crate::compaction::picker::Picker; + use crate::compaction::test_util::compaction_region_with_ssts; use crate::compaction::window::{WindowedCompactionPicker, file_time_bucket_span}; - use crate::region::options::{MergeMode, RegionOptions}; + use crate::region::options::{CompactionOptions, MergeMode, RegionOptions}; use crate::sst::file::{FileMeta, Level}; use crate::sst::file_purger::NoopFilePurger; use crate::sst::version::SstVersion; @@ -364,6 +374,29 @@ mod tests { } } + #[tokio::test] + async fn test_pick_output_file_size_threshold() { + let mut region = compaction_region_with_ssts( + [FileMeta { + file_id: FileId::random(), + time_range: (Timestamp::new_second(0), Timestamp::new_second(1)), + ..Default::default() + }], + Duration::from_secs(3600), + ) + .await; + let picker = WindowedCompactionPicker::new(Some(3600)); + + for (size, expected) in [(None, None), (Some(0), None), (Some(1024), Some(1024))] { + let CompactionOptions::Twcs(options) = &mut region.region_options.compaction; + options.max_output_file_size = size.map(ReadableSize); + + let output = picker.pick(®ion).await.unwrap().unwrap(); + assert_eq!(1, output.outputs.len()); + assert_eq!(expected, output.max_file_size); + } + } + #[test] fn test_pick_expired_ssts_without_marking_compacting() { let picker = WindowedCompactionPicker::new(None); diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index d2e03ac09a8..547d3f213f9 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -40,6 +40,7 @@ use tokio::sync::{Notify, Semaphore}; use crate::config::MitoConfig; use crate::engine::MitoEngine; use crate::engine::listener::{CompactionListener, EventListener}; +use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, column_metadata_to_column_schema, put_rows, }; @@ -137,6 +138,120 @@ async fn collect_stream_ts(stream: SendableRecordBatchStream) -> Vec { res } +#[rstest::rstest] +#[case::split("1B", true)] +#[case::unlimited("0B", false)] +#[tokio::test] +async fn test_strict_window_output_file_size( + #[values(false, true)] flat_format: bool, + #[case] max_output_file_size: &str, + #[case] expect_split: bool, +) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.max_output_file_size", max_output_file_size) + .insert_option("max_row_group_row_count", "50") + .build(); + let columns = crate::test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Each input SST spans both windows. Each series is large enough to exceed + // the tiny threshold, and must remain intact within its output window. + for key in ["a", "b", "c"] { + let rows = [0, 3600] + .into_iter() + .flat_map(|start| build_rows_for_key(key, start, start + 1000, start)) + .collect(); + put_rows( + &engine, + region_id, + Rows { + schema: columns.clone(), + rows, + }, + ) + .await; + flush(&engine, region_id).await; + } + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let before = RecordBatches::try_collect(stream).await.unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: api::v1::region::compact_request::Options::StrictWindow( + api::v1::region::StrictWindow { + window_seconds: 3600, + }, + ), + ..Default::default() + }), + ) + .await + .unwrap(); + + let version = engine.get_region(region_id).unwrap().version(); + assert!(version.ssts.levels()[0].files.is_empty()); + let files = version.ssts.levels()[1].files().collect::>(); + assert_eq!( + 6000, + files + .iter() + .map(|file| file.meta_ref().num_rows) + .sum::() + ); + for start in [0, 3600] { + let window_files = files + .iter() + .filter(|file| { + let (min, max) = file.time_range(); + min >= Timestamp::new_second(start) && max < Timestamp::new_second(start + 3600) + }) + .collect::>(); + if expect_split { + assert!(window_files.len() > 1, "each window must split by size"); + } else { + assert_eq!(1, window_files.len(), "zero must disable size splitting"); + } + let mut ranges = window_files + .iter() + .map(|file| file.primary_key_range().unwrap()) + .collect::>(); + ranges.sort_unstable(); + assert!(ranges.windows(2).all(|pair| pair[0].1 < pair[1].0)); + } + assert!(files.iter().all(|file| { + let (min, max) = file.time_range(); + min.value() / 3_600_000 == max.value() / 3_600_000 + })); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let after = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + sort_batches_and_print(&before, &["tag_0", "ts"]), + sort_batches_and_print(&after, &["tag_0", "ts"]) + ); +} + /// Flush may collapse versions within one file, but compaction must not promote /// them, even with no external overlaps or after rewriting a previous output. #[rstest::rstest] diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index fe603f9d1bc..89b413a5d8f 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -413,7 +413,9 @@ pub struct TwcsOptions { /// Compaction time window defined when creating tables. #[serde(with = "humantime_serde")] pub time_window: Option, - /// Compaction time window defined when creating tables. + /// Soft output SST size threshold for TWCS and strict-window compaction. + /// `None` or zero disables size-based splitting. With a primary key, files + /// split at series boundaries, so a single series may exceed the threshold. pub max_output_file_size: Option, /// Whether to use remote compaction. #[serde_as(as = "DisplayFromStr")] From 462172b9bbf54c22707a4bf1c9ff37e12151dfb5 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 07:46:40 +0000 Subject: [PATCH 24/63] perf(servers): reduce temporary memory in protocol handling (#9249) * perf(prom): release the remote write v1 decode buffer before writing remote_write_v1 kept the decoded builder alive until the handler returned, so the decompressed request payload stayed resident across the downstream write or pipeline await. With 8 concurrent large requests that is one extra copy of every payload held for the whole write. Rows and pipeline values own their data, so the builder can be dropped as soon as the conversion is done. Sustained-write A/B, 8 runs per side, 50M samples each: jemalloc allocated median drops 7.8% (155.0-162.7 MiB -> 141.4-158.6 MiB); samples per CPU second is unchanged (-0.6%, fully overlapping ranges). Signed-off-by: Dennis Zhuang * perf(servers): move row values into SQL JSON responses The format=json renderer cloned every serde_json::Value and kept the whole row set alive while building the response. Move the values out instead, so each row is released as soon as it is converted. Slicing the row to the schema width keeps the panic on rows narrower than the schema; a plain zip would silently truncate them. Duplicate column names still resolve to the last value and extra row values are still ignored. Isolated conversion measurements: live peak drops 33% on a 4096-row 16 KiB string fixture and 36% on a nested-JSON fixture, with no measured slowdown. Signed-off-by: Dennis Zhuang * perf(prom): release the compressed remote write body after decompression Both the v1 and v2 decoders held the compressed Bytes until they returned, which spans the whole protobuf decode and row conversion. Decompression copies the payload into an independent buffer, so the body can go as soon as it succeeds. The compression fallback, decode errors and request counting are unchanged. Signed-off-by: Dennis Zhuang * revert(prom): keep the remote write v1 decode buffer until the write finishes This reverts commit 8232c5b3bb00a656031620d1b30929e604d42717. The v1 decoder fabricates `&'static [u8]` pointing into its own decode buffer (prom_remote_write/types.rs), so the compiler checks nothing about that buffer's lifetime. Holding the builder until the handler returns is what keeps the decoder safe by construction; releasing it early made that safety depend on every consumer copying out of the buffer, which holds today but nothing enforces. Document the requirement at the binding instead. Signed-off-by: Dennis Zhuang * perf(prom): release the remote write v1 decode buffer before writing This reverts commit ca3af722d00ede48e259a798ee55c9320391c13e, restoring 8232c5b3bb00a656031620d1b30929e604d42717. The decoder borrows the decompressed buffer while parsing, but copies everything out when it builds rows: tag values through `PromValidationMode::decode_string`, column names through `to_owned`, and the only live borrows (`TableBuilder::col_indexes`) are dropped inside `as_insert_requests`. The resulting `ContextReq` holds prost types with no lifetime parameters, so it cannot reference the buffer. Record that at the binding so the next reader does not have to re-derive it from three files. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/servers/src/http/prom_store.rs | 12 ++++- src/servers/src/http/result/json_result.rs | 57 ++++++++++++++++++++-- src/servers/src/prom_remote_write/mod.rs | 2 + src/servers/src/prom_remote_write/v2.rs | 2 + 4 files changed, 66 insertions(+), 7 deletions(-) diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index cb7b6307c4b..60c000403fd 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -175,12 +175,20 @@ async fn remote_write_v1( processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def); } - let mut req = decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + let mut decoded = + decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + // Parsing borrows the decode buffer, but row building copies out of it: tag + // values through `decode_string`, column names through `to_owned`, and the + // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below + // references the buffer, so it need not span the write. let req = if processor.use_pipeline { + drop(decoded); processor.exec_pipeline().await? } else { - req.as_insert_requests() + let req = decoded.as_insert_requests(); + drop(decoded); + req }; let batches = into_prom_write_batches(req, query_ctx); diff --git a/src/servers/src/http/result/json_result.rs b/src/servers/src/http/result/json_result.rs index 537ea66cd02..74a75f26309 100644 --- a/src/servers/src/http/result/json_result.rs +++ b/src/servers/src/http/result/json_result.rs @@ -90,17 +90,20 @@ impl IntoResponse for JsonResponse { .to_string(), Some(GreptimeQueryOutput::Records(records)) => { - let schema = records.schema(); + // Borrow the schema field directly so the rows can be moved out. + let schema = &records.schema; let data: Vec> = records .rows - .iter() - .map(|row| { + .into_iter() + .map(|mut row| { + // Slicing keeps the out-of-bounds panic for short rows; + // a plain zip would silently truncate them. schema .column_schemas .iter() - .enumerate() - .map(|(i, col)| (col.name.clone(), row[i].clone())) + .zip(row[..schema.column_schemas.len()].iter_mut()) + .map(|(col, value)| (col.name.clone(), value.take())) .collect::>() }) .collect(); @@ -133,3 +136,47 @@ impl IntoResponse for JsonResponse { .into_response() } } + +#[cfg(test)] +mod tests { + use axum::body::to_bytes; + + use super::*; + + #[tokio::test] + async fn test_records_response_preserves_values_and_duplicate_columns() { + let response: JsonResponse = serde_json::from_value(json!({ + "output": [{"records": { + "schema": {"column_schemas": [ + {"name": "duplicate", "data_type": "String"}, + {"name": "nested", "data_type": "Json"}, + {"name": "duplicate", "data_type": "String"}, + {"name": "escaped\"column", "data_type": "String"} + ]}, + "rows": [ + ["discarded", {"array": [null, true, "中文"]}, "last", "line\n\\\""], + ["discarded", [1, {"key": "value"}], null, ""] + ] + }}], + "execution_time_ms": 7 + })) + .unwrap(); + + let response = response.into_response(); + assert_eq!(response.status(), axum::http::StatusCode::OK); + assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json"); + assert_eq!(response.headers()[&GREPTIME_DB_HEADER_FORMAT], "json"); + assert_eq!(response.headers()[&GREPTIME_DB_HEADER_EXECUTION_TIME], "7"); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert_eq!( + serde_json::from_slice::(&body).unwrap(), + json!({ + "data": [ + {"duplicate": "last", "nested": {"array": [null, true, "中文"]}, "escaped\"column": "line\n\\\""}, + {"duplicate": null, "nested": [1, {"key": "value"}], "escaped\"column": ""} + ], + "execution_time_ms": 7 + }) + ); + } +} diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index 73f968b6084..c9aabfb2745 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -77,6 +77,8 @@ pub fn decode_remote_write_request( // fallback to the other compression method try_decompress(!is_zstd, &body[..])? }; + // Decompression copied the payload out, so the compressed body is no longer needed. + drop(body); let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default); diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index b388071106f..2f1f4d410c4 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -143,6 +143,8 @@ pub(crate) fn decode_remote_write_v2( } else { try_decompress(!is_zstd, &body[..])? }; + // Decompression copied the payload out, so the compressed body is no longer needed. + drop(body); let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; drop(decode_timer); From fcb165cf776b51986093b1c873e9a7f14812290e Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Sun, 20 Sep 2026 08:53:00 +0000 Subject: [PATCH 25/63] feat(trace): support Jaeger queries for Trace V2 and optimize writes (follow-up to #9192) (#9257) feat(trace): support Jaeger queries for v2 and optimize fixed-column writes Signed-off-by: luofucong --- src/frontend/src/instance/jaeger.rs | 95 ++++++++-- src/servers/src/http/jaeger.rs | 48 +++-- src/servers/src/otlp/trace.rs | 1 + src/servers/src/otlp/trace/v2.rs | 261 ++++++++++++++++++---------- src/servers/src/row_writer.rs | 21 +-- tests-integration/tests/http.rs | 188 ++++++++++++++++++++ 6 files changed, 481 insertions(+), 133 deletions(-) diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index 4d3b137f1c9..556f7640de8 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -23,16 +23,18 @@ use common_catalog::consts::{ }; use common_function::function::FunctionRef; use common_function::scalars::json::json_get::{ - JsonGetBool, JsonGetFloat, JsonGetInt, JsonGetString, + JsonGetBool, JsonGetFloat, JsonGetInt, JsonGetString, JsonGetWithType, }; use common_function::scalars::udf::create_udf; use common_query::{Output, OutputData}; use common_recordbatch::adapter::RecordBatchStreamAdapter; use common_recordbatch::util; use common_telemetry::warn; +use datafusion::common::ScalarValue; use datafusion::dataframe::DataFrame; use datafusion::execution::SessionStateBuilder; use datafusion::execution::context::SessionContext; +use datafusion::functions::core::expr_fn::coalesce; use datafusion::functions_window::expr_fn::row_number; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{Expr, ExprFunctionExt, SortExpr, col, lit, lit_timestamp_nano, wildcard}; @@ -43,9 +45,9 @@ use servers::error::{ }; use servers::http::jaeger::{JAEGER_QUERY_TABLE_NAME_KEY, QueryTraceParams, TraceUserAgent}; use servers::otlp::trace::{ - DURATION_NANO_COLUMN, KEY_OTEL_STATUS_ERROR_KEY, SERVICE_NAME_COLUMN, SPAN_ATTRIBUTES_COLUMN, - SPAN_KIND_COLUMN, SPAN_KIND_PREFIX, SPAN_NAME_COLUMN, SPAN_STATUS_CODE, SPAN_STATUS_ERROR, - TIMESTAMP_COLUMN, TRACE_ID_COLUMN, + DURATION_NANO_COLUMN, KEY_OTEL_STATUS_ERROR_KEY, RESOURCE_ATTRIBUTES_COLUMN, + SERVICE_NAME_COLUMN, SPAN_ATTRIBUTES_COLUMN, SPAN_KIND_COLUMN, SPAN_KIND_PREFIX, + SPAN_NAME_COLUMN, SPAN_STATUS_CODE, SPAN_STATUS_ERROR, TIMESTAMP_COLUMN, TRACE_ID_COLUMN, }; use servers::query_handler::JaegerQueryHandler; use session::context::QueryContextRef; @@ -367,11 +369,16 @@ async fn query_trace_table( schema: ctx.current_schema(), })?; - let is_data_model_v1 = table::requests::is_trace_v1_table(&table.table_info()); + let table_info = table.table_info(); + let data_model = table_info + .meta + .options + .extra_options + .get(table::requests::TABLE_DATA_MODEL) + .map(String::as_str); // collect to set - let col_names = table - .table_info() + let col_names = table_info .meta .field_column_names() .map(|s| format!("\"{}\"", s)) @@ -389,7 +396,7 @@ async fn query_trace_table( let dataframe = filters .into_iter() .chain(tags.map_or(Ok(vec![]), |t| { - tags_filters(&dataframe, t, is_data_model_v1, &col_names) + tags_filters(&dataframe, t, data_model, &col_names) })?) .try_fold(dataframe, |df, expr| { df.filter(expr).context(DataFusionSnafu) @@ -518,8 +525,9 @@ fn create_df_context(query_engine: &QueryEngineRef) -> ServerResult = vec![ + Arc::new(JsonGetWithType::default()), Arc::new(JsonGetInt::default()), Arc::new(JsonGetFloat::default()), Arc::new(JsonGetBool::default()), @@ -600,6 +608,65 @@ fn json_tag_filters( Ok(filters) } +/// Resolves each tag per row, falling back to Resource when the typed Span value is null. +fn json2_tag_filters( + dataframe: &DataFrame, + tags: HashMap, +) -> ServerResult> { + let get = dataframe + .registry() + .udf(JsonGetWithType::NAME) + .context(DataFusionSnafu)?; + tags.into_iter() + .map(|(key, value)| { + if key == KEY_OTEL_STATUS_ERROR_KEY && value == JsonValue::Bool(true) { + return Ok(col(SPAN_STATUS_CODE).eq(lit(SPAN_STATUS_ERROR))); + } + // JSON quoting preserves literal dots, quotes, and backslashes in attribute keys. + let path = lit(format!("$[{}]", JsonValue::String(key))); + let (value, value_type) = match value { + JsonValue::String(value) => ( + ScalarValue::Utf8View(Some(value)), + ScalarValue::Utf8View(None), + ), + JsonValue::Bool(value) => ( + ScalarValue::Boolean(Some(value)), + ScalarValue::Boolean(None), + ), + JsonValue::Number(value) => { + if let Some(value) = value.as_i64() { + (ScalarValue::Int64(Some(value)), ScalarValue::Int64(None)) + } else if let Some(value) = value.as_u64() { + (ScalarValue::UInt64(Some(value)), ScalarValue::UInt64(None)) + } else if let Some(value) = value.as_f64() { + ( + ScalarValue::Float64(Some(value)), + ScalarValue::Float64(None), + ) + } else { + return Ok(lit(false)); + } + } + JsonValue::Null => (ScalarValue::Utf8View(None), ScalarValue::Utf8View(None)), + JsonValue::Array(_) | JsonValue::Object(_) => return Ok(lit(false)), + }; + let attribute = coalesce(vec![ + get.call(vec![ + col(SPAN_ATTRIBUTES_COLUMN), + path.clone(), + lit(value_type.clone()), + ]), + get.call(vec![col(RESOURCE_ATTRIBUTES_COLUMN), path, lit(value_type)]), + ]); + Ok(if value.is_null() { + attribute.is_null() + } else { + attribute.eq(lit(value)) + }) + }) + .collect() +} + /// Helper function to check if span_key or resource_key exists in col_names and create an expression. /// If neither exists, logs a warning and returns None. #[inline] @@ -679,13 +746,13 @@ fn flatten_tag_filters( fn tags_filters( dataframe: &DataFrame, tags: HashMap, - is_data_model_v1: bool, + data_model: Option<&str>, col_names: &HashSet, ) -> ServerResult> { - if is_data_model_v1 { - flatten_tag_filters(tags, col_names) - } else { - json_tag_filters(dataframe, tags) + match data_model { + Some(table::requests::TABLE_DATA_MODEL_TRACE_V1) => flatten_tag_filters(tags, col_names), + Some(table::requests::TABLE_DATA_MODEL_TRACE_V2) => json2_tag_filters(dataframe, tags), + _ => json_tag_filters(dataframe, tags), } } diff --git a/src/servers/src/http/jaeger.rs b/src/servers/src/http/jaeger.rs index db843bf74e2..32ef6a61682 100644 --- a/src/servers/src/http/jaeger.rs +++ b/src/servers/src/http/jaeger.rs @@ -45,9 +45,9 @@ use crate::otlp::trace::{ KEY_OTEL_STATUS_ERROR_KEY, KEY_OTEL_STATUS_MESSAGE, KEY_OTEL_TRACE_STATE, KEY_SERVICE_NAME, KEY_SPAN_KIND, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_NAME_COLUMN, SCOPE_VERSION_COLUMN, SERVICE_NAME_COLUMN, SPAN_ATTRIBUTES_COLUMN, SPAN_EVENTS_COLUMN, SPAN_ID_COLUMN, - SPAN_KIND_COLUMN, SPAN_KIND_PREFIX, SPAN_NAME_COLUMN, SPAN_STATUS_CODE, SPAN_STATUS_ERROR, - SPAN_STATUS_MESSAGE_COLUMN, SPAN_STATUS_PREFIX, SPAN_STATUS_UNSET, TIMESTAMP_COLUMN, - TRACE_ID_COLUMN, TRACE_STATE_COLUMN, + SPAN_KIND_COLUMN, SPAN_KIND_PREFIX, SPAN_LINKS_COLUMN, SPAN_NAME_COLUMN, SPAN_STATUS_CODE, + SPAN_STATUS_ERROR, SPAN_STATUS_MESSAGE_COLUMN, SPAN_STATUS_PREFIX, SPAN_STATUS_UNSET, + TIMESTAMP_COLUMN, TRACE_ID_COLUMN, TRACE_STATE_COLUMN, }; use crate::query_handler::JaegerQueryHandlerRef; @@ -779,6 +779,7 @@ fn traces_from_records(records: HttpRecordsOutput) -> Result> { for row in records.rows.into_iter() { let mut span = Span::default(); let mut service_name = None; + let mut parent_span_id = None; let mut resource_tags = vec![]; for (idx, cell) in row.into_iter().enumerate() { @@ -818,14 +819,14 @@ fn traces_from_records(records: HttpRecordsOutput) -> Result> { } } SPAN_ATTRIBUTES_COLUMN => { - // for v0 data model, span_attributes are nested as a json + // For v0 and v2, span_attributes are nested as a JSON // data structure if let JsonValue::Object(span_attrs) = cell { span.tags.extend(object_to_tags(span_attrs)); } } RESOURCE_ATTRIBUTES_COLUMN => { - // for v0 data model, resource_attributes are nested as a json + // For v0 and v2, resource_attributes are nested as a JSON // data structure if let JsonValue::Object(mut resource_attrs) = cell { @@ -834,14 +835,26 @@ fn traces_from_records(records: HttpRecordsOutput) -> Result> { } } PARENT_SPAN_ID_COLUMN => { - if let JsonValue::String(parent_span_id) = cell - && !parent_span_id.is_empty() + if let JsonValue::String(id) = cell + && !id.is_empty() { - span.references.push(Reference { - trace_id: span.trace_id.clone(), - span_id: parent_span_id, - ref_type: REF_TYPE_CHILD_OF.to_string(), - }); + parent_span_id = Some(id); + } + } + SPAN_LINKS_COLUMN => { + if let JsonValue::Array(links) = cell { + for link in links { + if let (Some(trace_id), Some(span_id)) = ( + link.get("trace_id").and_then(JsonValue::as_str), + link.get("span_id").and_then(JsonValue::as_str), + ) { + span.references.push(Reference { + trace_id: trace_id.to_string(), + span_id: span_id.to_string(), + ref_type: "FOLLOWS_FROM".to_string(), + }); + } + } } } SPAN_EVENTS_COLUMN => { @@ -994,6 +1007,17 @@ fn traces_from_records(records: HttpRecordsOutput) -> Result> { } } + if let Some(parent_span_id) = parent_span_id { + span.references.insert( + 0, + Reference { + trace_id: span.trace_id.clone(), + span_id: parent_span_id, + ref_type: REF_TYPE_CHILD_OF.to_string(), + }, + ); + } + if let Some(service_name) = service_name { if !service_to_resource_attributes.contains_key(&service_name) { service_to_resource_attributes.insert(service_name.clone(), resource_tags); diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index ca1fe1defd4..60865472086 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -39,6 +39,7 @@ pub const TIMESTAMP_END_COLUMN: &str = "timestamp_end"; pub const SPAN_STATUS_MESSAGE_COLUMN: &str = "span_status_message"; pub const SPAN_ATTRIBUTES_COLUMN: &str = "span_attributes"; pub const SPAN_EVENTS_COLUMN: &str = "span_events"; +pub const SPAN_LINKS_COLUMN: &str = "span_links"; pub const SCOPE_NAME_COLUMN: &str = "scope_name"; pub const SCOPE_VERSION_COLUMN: &str = "scope_version"; /// Column containing instrumentation scope attributes as a JSON2 object. diff --git a/src/servers/src/otlp/trace/v2.rs b/src/servers/src/otlp/trace/v2.rs index d4628c3b884..dd43037822f 100644 --- a/src/servers/src/otlp/trace/v2.rs +++ b/src/servers/src/otlp/trace/v2.rs @@ -14,7 +14,6 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, ColumnSchema, RowInsertRequests, SemanticType}; -use common_grpc::precision::Precision; use snafu::ensure; use crate::error::{Result, TimestampOverflowSnafu}; @@ -23,11 +22,10 @@ use crate::otlp::trace::v1::span_duration_nano; use crate::otlp::trace::{ DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, SCOPE_NAME_COLUMN, SCOPE_VERSION_COLUMN, SERVICE_NAME_COLUMN, - SPAN_ATTRIBUTES_COLUMN, SPAN_EVENTS_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_COLUMN, SPAN_NAME_COLUMN, - SPAN_STATUS_CODE, SPAN_STATUS_MESSAGE_COLUMN, TIMESTAMP_COLUMN, TIMESTAMP_END_COLUMN, - TRACE_ID_COLUMN, TRACE_STATE_COLUMN, + SPAN_ATTRIBUTES_COLUMN, SPAN_EVENTS_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_COLUMN, + SPAN_LINKS_COLUMN, SPAN_NAME_COLUMN, SPAN_STATUS_CODE, SPAN_STATUS_MESSAGE_COLUMN, + TIMESTAMP_COLUMN, TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_STATE_COLUMN, }; -use crate::otlp::utils::{make_column_data, make_string_column_data}; use crate::row_writer::{self, MultiTableData, TableData}; // Preallocate for the fixed v2 schema: 3 timing columns, 3 span/trace IDs, @@ -35,6 +33,79 @@ use crate::row_writer::{self, MultiTableData, TableData}; // columns, 1 service name, 3 JSON2 attribute columns, and 2 JSON event/link columns. const APPROXIMATE_COLUMN_COUNT: usize = 19; +struct FixedTraceColumnIndexes { + timestamp: usize, + timestamp_end: usize, + duration_nano: usize, + parent_span_id: usize, + trace_id: usize, + span_id: usize, + span_kind: usize, + span_name: usize, + span_status_code: usize, + span_status_message: usize, + trace_state: usize, + scope_name: usize, + scope_version: usize, + service_name: usize, + span_attributes: usize, + scope_attributes: usize, + resource_attributes: usize, + span_events: usize, + span_links: usize, +} + +impl FixedTraceColumnIndexes { + fn resolve(writer: &mut TableData) -> Result { + let timestamp = writer.ensure_column(api::v1::helper::time_index_column_schema( + TIMESTAMP_COLUMN, + ColumnDataType::TimestampNanosecond, + ))?; + let mut field = |name: &str, datatype: ColumnDataType| { + writer.ensure_column(ColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }) + }; + Ok(Self { + timestamp, + timestamp_end: field(TIMESTAMP_END_COLUMN, ColumnDataType::TimestampNanosecond)?, + duration_nano: field(DURATION_NANO_COLUMN, ColumnDataType::Int64)?, + parent_span_id: field(PARENT_SPAN_ID_COLUMN, ColumnDataType::String)?, + trace_id: field(TRACE_ID_COLUMN, ColumnDataType::String)?, + span_id: field(SPAN_ID_COLUMN, ColumnDataType::String)?, + span_kind: field(SPAN_KIND_COLUMN, ColumnDataType::String)?, + span_name: field(SPAN_NAME_COLUMN, ColumnDataType::String)?, + span_status_code: field(SPAN_STATUS_CODE, ColumnDataType::String)?, + span_status_message: field(SPAN_STATUS_MESSAGE_COLUMN, ColumnDataType::String)?, + trace_state: field(TRACE_STATE_COLUMN, ColumnDataType::String)?, + scope_name: field(SCOPE_NAME_COLUMN, ColumnDataType::String)?, + scope_version: field(SCOPE_VERSION_COLUMN, ColumnDataType::String)?, + service_name: writer.ensure_column(ColumnSchema { + column_name: SERVICE_NAME_COLUMN.to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + })?, + span_attributes: writer.ensure_column(row_writer::build_json2_column_schema( + SPAN_ATTRIBUTES_COLUMN, + ))?, + scope_attributes: writer.ensure_column(row_writer::build_json2_column_schema( + SCOPE_ATTRIBUTES_COLUMN, + ))?, + resource_attributes: writer.ensure_column(row_writer::build_json2_column_schema( + RESOURCE_ATTRIBUTES_COLUMN, + ))?, + span_events: writer + .ensure_column(row_writer::build_json_column_schema(SPAN_EVENTS_COLUMN))?, + span_links: writer + .ensure_column(row_writer::build_json_column_schema(SPAN_LINKS_COLUMN))?, + }) + } +} + /// Converts trace spans into row insert requests for the main v2 trace table. pub(super) fn v2_to_grpc_main_insert_requests( spans: &[TraceSpan], @@ -48,13 +119,21 @@ pub(super) fn v2_to_grpc_main_insert_requests( /// Builds the fixed row-oriented payload for the main v2 trace table. fn build_trace_table_data(spans: &[TraceSpan]) -> Result { let mut writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); + if spans.is_empty() { + return Ok(writer); + } + let columns = FixedTraceColumnIndexes::resolve(&mut writer)?; for span in spans { - write_span_to_row(&mut writer, span)?; + write_span_to_row(&mut writer, span, &columns)?; } Ok(writer) } -fn write_span_to_row(writer: &mut TableData, span: &TraceSpan) -> Result<()> { +fn write_span_to_row( + writer: &mut TableData, + span: &TraceSpan, + columns: &FixedTraceColumnIndexes, +) -> Result<()> { ensure!( span.start_in_nanosecond <= i64::MAX as u64, TimestampOverflowSnafu { @@ -69,90 +148,92 @@ fn write_span_to_row(writer: &mut TableData, span: &TraceSpan) -> Result<()> { ); let mut row = writer.alloc_one_row(); - row_writer::write_ts_to_nanos( - writer, - TIMESTAMP_COLUMN, - Some(span.start_in_nanosecond as i64), - Precision::Nanosecond, - &mut row, - )?; - row_writer::write_fields( - writer, - vec![ - make_column_data( - TIMESTAMP_END_COLUMN, - ColumnDataType::TimestampNanosecond, - Some(ValueData::TimestampNanosecondValue( - span.end_in_nanosecond as i64, - )), - ), - make_column_data( - DURATION_NANO_COLUMN, - ColumnDataType::Int64, - Some(ValueData::I64Value(span_duration_nano(span))), - ), - make_string_column_data(PARENT_SPAN_ID_COLUMN, span.parent_span_id.clone()), - make_string_column_data(TRACE_ID_COLUMN, Some(span.trace_id.clone())), - make_string_column_data(SPAN_ID_COLUMN, Some(span.span_id.clone())), - make_string_column_data(SPAN_KIND_COLUMN, Some(span.span_kind.clone())), - make_string_column_data(SPAN_NAME_COLUMN, Some(span.span_name.clone())), - make_string_column_data(SPAN_STATUS_CODE, Some(span.span_status_code.clone())), - make_string_column_data( - SPAN_STATUS_MESSAGE_COLUMN, - Some(span.span_status_message.clone()), - ), - make_string_column_data(TRACE_STATE_COLUMN, Some(span.trace_state.clone())), - make_string_column_data(SCOPE_NAME_COLUMN, Some(span.scope_name.clone())), - make_string_column_data(SCOPE_VERSION_COLUMN, Some(span.scope_version.clone())), - ] - .into_iter(), - &mut row, - )?; - row_writer::write_by_schema( - writer, - std::iter::once(( - ColumnSchema { - column_name: SERVICE_NAME_COLUMN.to_string(), - datatype: ColumnDataType::String as i32, - semantic_type: SemanticType::Tag as i32, - ..Default::default() - }, + let duration = span_duration_nano(span); + for (index, value) in [ + ( + columns.timestamp, + Some(ValueData::TimestampNanosecondValue( + span.start_in_nanosecond as i64, + )), + ), + ( + columns.timestamp_end, + Some(ValueData::TimestampNanosecondValue( + span.end_in_nanosecond as i64, + )), + ), + (columns.duration_nano, Some(ValueData::I64Value(duration))), + ( + columns.parent_span_id, + span.parent_span_id.clone().map(ValueData::StringValue), + ), + ( + columns.trace_id, + Some(ValueData::StringValue(span.trace_id.clone())), + ), + ( + columns.span_id, + Some(ValueData::StringValue(span.span_id.clone())), + ), + ( + columns.span_kind, + Some(ValueData::StringValue(span.span_kind.clone())), + ), + ( + columns.span_name, + Some(ValueData::StringValue(span.span_name.clone())), + ), + ( + columns.span_status_code, + Some(ValueData::StringValue(span.span_status_code.clone())), + ), + ( + columns.span_status_message, + Some(ValueData::StringValue(span.span_status_message.clone())), + ), + ( + columns.trace_state, + Some(ValueData::StringValue(span.trace_state.clone())), + ), + ( + columns.scope_name, + Some(ValueData::StringValue(span.scope_name.clone())), + ), + ( + columns.scope_version, + Some(ValueData::StringValue(span.scope_version.clone())), + ), + ( + columns.service_name, span.service_name.clone().map(ValueData::StringValue), - )), - &mut row, - )?; - - row_writer::write_json2( - writer, - SPAN_ATTRIBUTES_COLUMN, - &span.span_attributes, - &mut row, - )?; - row_writer::write_json2( - writer, - SCOPE_ATTRIBUTES_COLUMN, - span.scope_attributes.as_ref(), - &mut row, - )?; - row_writer::write_json2( - writer, - RESOURCE_ATTRIBUTES_COLUMN, - span.resource_attributes.as_ref(), - &mut row, - )?; - - row_writer::write_json( - writer, - SPAN_EVENTS_COLUMN, - span.span_events.clone().into(), - &mut row, - )?; - row_writer::write_json( - writer, - "span_links", - span.span_links.clone().into(), - &mut row, - )?; + ), + ( + columns.span_attributes, + Some(row_writer::encode_json2(&span.span_attributes)?), + ), + ( + columns.scope_attributes, + Some(row_writer::encode_json2(span.scope_attributes.as_ref())?), + ), + ( + columns.resource_attributes, + Some(row_writer::encode_json2(span.resource_attributes.as_ref())?), + ), + ( + columns.span_events, + Some(ValueData::BinaryValue( + jsonb::Value::from(span.span_events.clone()).to_vec(), + )), + ), + ( + columns.span_links, + Some(ValueData::BinaryValue( + jsonb::Value::from(span.span_links.clone()).to_vec(), + )), + ), + ] { + row[index].value_data = value; + } writer.add_row(row); Ok(()) diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 7fce0d60711..0661e733510 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -327,7 +327,7 @@ pub fn write_json( ) } -fn build_json2_column_schema(name: impl ToString) -> ColumnSchema { +pub(crate) fn build_json2_column_schema(name: impl ToString) -> ColumnSchema { let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::new_v2()))); let mut options = ColumnOptions::default(); options.options.insert( @@ -349,13 +349,8 @@ fn build_json2_column_schema(name: impl ToString) -> ColumnSchema { } } -/// Writes a serializable value into a JSON2 field using the default v2 layout. -pub(crate) fn write_json2( - table_data: &mut TableData, - name: impl ToString, - value: impl Serialize, - one_row: &mut Vec, -) -> Result<()> { +/// Encodes a JSON2 field without constructing its column schema. +pub(crate) fn encode_json2(value: impl Serialize) -> Result { let json = serde_json::to_value(value).context(ToJsonSnafu)?; let value = JsonSettings::new_v2() .encode(json) @@ -366,15 +361,7 @@ pub(crate) fn write_json2( } .fail(); }; - - write_by_schema( - table_data, - std::iter::once(( - build_json2_column_schema(name), - Some(ValueData::JsonValue(encode_json_value(*value))), - )), - one_row, - ) + Ok(ValueData::JsonValue(encode_json_value(*value))) } pub(crate) fn write_by_schema( diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index cd614933353..1ad39ae1726 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -194,6 +194,7 @@ macro_rules! http_tests { test_log_query, test_jaeger_query_api, test_jaeger_query_api_for_trace_v1, + test_jaeger_query_api_for_trace_v2, test_influxdb_write, test_influxdb_write_with_hints, @@ -10406,6 +10407,193 @@ pub async fn test_jaeger_query_api(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_jaeger_query_api_for_trace_v2( + store_type: StorageType, +) -> Result<(), Box> { + common_telemetry::init_default_ut_logging(); + + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_jaeger_query_api_v2").await; + let client = TestClient::new(app).await; + let request: ExportTraceServiceRequest = serde_json::from_value(json!({ + "resourceSpans": [{ + "resource": {"attributes": [ + make_string_attr("service.name", "jaeger-v2"), + make_string_attr("region", "west"), + make_string_attr("shared", "resource"), + make_string_attr("nullable", "resource"), + make_int_attr("count", 7), + make_int_attr("complex", 7) + ]}, + "scopeSpans": [{ + "scope": {"name": "v2-scope", "version": "1.0"}, + "spans": [{ + "traceId": "c05d7a4ec8e1f231f02ed6e8da8655b4", + "spanId": "9630f2916e2f7909", "parentSpanId": "8f847259b0f6e1ab", + "name": "op", "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000", + "attributes": [ + make_string_attr("shared", "span"), + {"key": "nullable"}, + make_string_attr("a\"b\\c.d", "escaped"), + make_int_attr("http.status_code", 500), + make_bool_attr("enabled", true), + {"key": "ratio", "value": {"doubleValue": 1.5}}, + make_string_attr("count", "invalid"), + {"key": "complex", "value": {"arrayValue": {"values": [{"intValue":"7"}]}}} + ], + "events": [{"name": "cache.hit", "timeUnixNano": "1736480942444400000", + "attributes": [make_int_attr("event.code", 7)]}], + "links": [{"traceId": "cc9e0991a2e63d274984bd44ee669203", + "spanId": "8f847259b0f6e1ab"}], + "status": {"code": 2, "message": "failure"} + }, { + "traceId": "cc9e0991a2e63d274984bd44ee669203", + "spanId": "8f847259b0f6e1ab", "name": "op", "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000" + }] + }] + }] + }))?; + let response = send_req( + &client, + vec![ + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ), + ( + HeaderName::from_static("x-greptime-pipeline-name"), + HeaderValue::from_static(GREPTIME_INTERNAL_TRACE_PIPELINE_V2_NAME), + ), + ( + HeaderName::from_static("x-greptime-trace-table-name"), + HeaderValue::from_static("jaeger_v2"), + ), + ], + "/v1/otlp/v1/traces", + request.encode_to_vec(), + false, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + + for (path, expected) in [ + ("/v1/jaeger/api/services", json!(["jaeger-v2"])), + ( + "/v1/jaeger/api/operations?service=jaeger-v2", + json!([{"name":"op", "spanKind":"server"}]), + ), + ( + "/v1/jaeger/api/services/jaeger-v2/operations", + json!(["op"]), + ), + ] { + let response = client + .get(path) + .header("x-greptime-trace-table-name", "jaeger_v2") + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!(body["data"], expected, "{path}: {body}"); + } + + for flushed in [false, true] { + if flushed { + let response = client + .post("/v1/sql") + .form(&[("sql", "ADMIN FLUSH_TABLE('jaeger_v2')")]) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + } + for (tags, expected) in [ + (json!({"shared":"span"}), 1), + (json!({"shared":"resource"}), 1), + (json!({"region":"west"}), 2), + (json!({"nullable":"resource"}), 2), + (json!({"count":7}), 2), + (json!({"complex":7}), 2), + (json!({"a\"b\\c.d":"escaped"}), 1), + ( + json!({"http.status_code":500, "enabled":true, "ratio":1.5}), + 1, + ), + (json!({"error":true}), 1), + (json!({"missing":"absent"}), 0), + (json!({"missing":null}), 2), + (json!({"count":u64::MAX}), 0), + ] { + let url = format!( + "/v1/jaeger/api/traces?service=jaeger-v2&start=1736480942444376&end=1736480942444500&tags={}", + encode(&tags.to_string()) + ); + let response = client + .get(&url) + .header("user-agent", if flushed { "Grafana" } else { "Jaeger" }) + .header("x-greptime-trace-table-name", "jaeger_v2") + .send() + .await; + let status = response.status(); + let body: Value = serde_json::from_str(&response.text().await)?; + assert_eq!(status, StatusCode::OK, "{tags}: {body}"); + assert_eq!( + body["data"].as_array().map_or(0, Vec::len), + expected, + "{tags}: {body}" + ); + } + let response = client + .get("/v1/jaeger/api/traces/c05d7a4ec8e1f231f02ed6e8da8655b4") + .header("x-greptime-trace-table-name", "jaeger_v2") + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = serde_json::from_str(&response.text().await)?; + let process = &body["data"][0]["processes"]["p1"]; + assert_eq!(process["serviceName"], "jaeger-v2"); + assert!( + process["tags"] + .as_array() + .ok_or("Expected process tags")? + .iter() + .any(|tag| tag["key"] == "region" && tag["value"] == "west") + ); + let span = &body["data"][0]["spans"][0]; + assert_eq!(span["startTime"], 1736480942444376u64); + assert_eq!(span["duration"], 123); + assert_eq!( + span["references"], + json!([ + {"refType":"CHILD_OF", "traceID":"c05d7a4ec8e1f231f02ed6e8da8655b4", "spanID":"8f847259b0f6e1ab"}, + {"refType":"FOLLOWS_FROM", "traceID":"cc9e0991a2e63d274984bd44ee669203", "spanID":"8f847259b0f6e1ab"}, + ]) + ); + assert_eq!(span["logs"][0]["timestamp"], 1736480942444400u64); + assert!( + span["logs"][0]["fields"] + .as_array() + .ok_or("Expected event fields")? + .iter() + .any(|field| field["key"] == "event.code" && field["value"] == 7) + ); + assert!( + span["tags"] + .as_array() + .ok_or("Expected span tags")? + .iter() + .any(|tag| tag["key"] == "http.status_code" + && tag["type"] == "int64" + && tag["value"] == 500) + ); + } + guard.remove_all().await; + Ok(()) +} + pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From e9dd79d1d2a9357ae712219976100b173f1dddba Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Sun, 20 Sep 2026 14:30:12 +0000 Subject: [PATCH 26/63] fix(mito2): compare primary key ranges across schema versions (#9205) * fix(mito2): compare primary key ranges across schema versions Bind FileHandle ranges to the pinned region schema and append cached constant defaults to historical Dense keys. Preserve raw SST statistics, reject inexact bounds, and avoid invalidating views for unrelated metadata changes. Cover schema evolution, default changes, and tombstone retention through real compaction and reopen regressions. Signed-off-by: Lei, HUANG * refactor(mito2): report invalid primary key ranges Signed-off-by: Lei, HUANG * refactor(mito2): scope primary key ranges to comparisons Keep raw PK bounds only in FileHandleInner and move schema-aware mapping and caching into task-local comparison contexts. Use explicit contexts for compaction overlap checks, window aggregation, and series scans. Preserve pinned-schema isolation, late statistics, and shared file lifecycle state without rebinding every handle. Cover cache isolation across region owners and adapt range fixtures to real Dense encodings. All 1530 mito2 tests and Clippy for all targets with the testing feature pass. Signed-off-by: Lei, HUANG * refactor(mito2): cache aligned primary key ranges per file Replace task-local range maps with a single-slot cache in FileHandleInner, keyed by the target schema version. Preserve raw bounds for realignment across snapshots and default changes. Share schema mappers across comparison paths and use copy-on-write SST lists for metadata updates. Simplify range mapping to accept encoded bounds and assert the same-table contract at the file accessor. Signed-off-by: Lei, HUANG * docs(mito2): clarify primary key mapper schema snapshot Signed-off-by: Lei, HUANG * test(mito2): align primary key range fixtures with table contract Remove obsolete cross-table fallback expectations after region validation became a caller contract. Give compaction fixtures matching table identities, including the active-window L1 scenario. Clarify the mapper precondition and format the simplified alignment call. All 1530 mito2 tests and Clippy for all targets with the testing feature pass. Signed-off-by: Lei, HUANG * test(mito2): enable filesystem GC in release unit tests Let unit tests use the filesystem-backed object-store GC path regardless of optimization profile. Keep the production release GC selection unchanged. Signed-off-by: Lei, HUANG * perf(mito-codec): skip release value decoding in PK prefix counts Validate field values only in debug builds and unit tests while keeping boundary, truncation, and trailing-byte checks in every build. Cover the linked library in debug and release integration tests, and verify that release unit tests still perform value validation. Signed-off-by: Lei, HUANG * test(mito-codec): remove redundant prefix integration tests Retain the codec unit tests and cross-schema compaction regressions while dropping the standalone build-profile test file and its release-only invalid-value expectation. Signed-off-by: Lei, HUANG * ci: wait for MySQL to accept authenticated TCP queries Add a healthcheck using the configured test account and database. Docker Compose --wait previously only observed container startup because the fixture image had no healthcheck, allowing metasrv to connect before MySQL initialization completed. Verify readiness with SELECT 1 over TCP rather than the initialization socket. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito-codec/src/error.rs | 11 +- src/mito-codec/src/row_converter/dense.rs | 224 +++-- src/mito2/benches/bench_compaction_picker.rs | 9 +- src/mito2/src/compaction/compactor.rs | 6 +- src/mito2/src/compaction/last_non_null.rs | 18 +- src/mito2/src/compaction/overlap.rs | 51 +- src/mito2/src/compaction/picker.rs | 3 +- src/mito2/src/compaction/run.rs | 98 +- src/mito2/src/compaction/test_util.rs | 27 +- src/mito2/src/compaction/twcs.rs | 280 +++++- src/mito2/src/compaction/window.rs | 2 +- src/mito2/src/engine/compaction_test.rs | 321 ++++++- src/mito2/src/error.rs | 20 +- src/mito2/src/read/scan_region.rs | 18 +- src/mito2/src/read/series_reader.rs | 2 +- src/mito2/src/region/opener.rs | 10 +- src/mito2/src/region/version.rs | 9 +- src/mito2/src/sst.rs | 1 + src/mito2/src/sst/file.rs | 104 ++- src/mito2/src/sst/file_purger.rs | 8 +- src/mito2/src/sst/parquet/metadata.rs | 15 + src/mito2/src/sst/primary_key.rs | 847 ++++++++++++++++++ src/mito2/src/sst/version.rs | 43 +- tests-integration/fixtures/docker-compose.yml | 13 + 24 files changed, 1929 insertions(+), 211 deletions(-) create mode 100644 src/mito2/src/sst/primary_key.rs diff --git a/src/mito-codec/src/error.rs b/src/mito-codec/src/error.rs index 3d3b0efcc07..b266ea9b04f 100644 --- a/src/mito-codec/src/error.rs +++ b/src/mito-codec/src/error.rs @@ -74,6 +74,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid dense primary key: {}", reason))] + InvalidDensePrimaryKey { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Encode null value"))] IndexEncodeNull { #[snafu(implicit)] @@ -101,7 +108,9 @@ impl ErrorExt for Error { StatusCode::InvalidArguments } NotSupportedField { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, - InvalidSparsePrimaryKey { .. } => StatusCode::InvalidArguments, + InvalidSparsePrimaryKey { .. } | InvalidDensePrimaryKey { .. } => { + StatusCode::InvalidArguments + } EvaluateFilter { source, .. } => source.status_code(), } } diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs index 6cc70feaeaf..56289a08b11 100644 --- a/src/mito-codec/src/row_converter/dense.rs +++ b/src/mito-codec/src/row_converter/dense.rs @@ -290,65 +290,84 @@ impl SortField { bytes: &[u8], deserializer: &mut Deserializer<&[u8]>, ) -> Result { - let pos = deserializer.position(); - if bytes[pos] == 0 { - deserializer.advance(1); - return Ok(1); - } - - Self::skip_deserialize_by_type(self.encode_data_type(), bytes, deserializer) + let len = self.encoded_field_len(&bytes[deserializer.position()..])?; + deserializer.advance(len); + Ok(len) } - fn skip_deserialize_by_type( - data_type: &ConcreteDataType, - bytes: &[u8], - deserializer: &mut Deserializer<&[u8]>, - ) -> Result { - let to_skip = match data_type { - ConcreteDataType::Boolean(_) => 2, - ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2, + /// Checks field boundaries before any unchecked reads in memcomparable. + fn encoded_field_len(&self, bytes: &[u8]) -> Result { + let invalid = || { + error::InvalidDensePrimaryKeySnafu { + reason: "truncated field or invalid encoding", + } + .build() + }; + match bytes.first() { + Some(0) => return Ok(1), + Some(1) => {} + _ => return Err(invalid()), + } + let len = match self.encode_data_type() { + ConcreteDataType::Boolean(_) + | ConcreteDataType::Int8(_) + | ConcreteDataType::UInt8(_) => 2, ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3, - ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5, - ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9, - ConcreteDataType::Float32(_) => 5, - ConcreteDataType::Float64(_) => 9, - ConcreteDataType::Binary(_) - | ConcreteDataType::Json(_) - | ConcreteDataType::Vector(_) => { - // Now the encoder encode binary as a list of bytes so we can't use - // skip bytes. - let pos_before = deserializer.position(); - let mut current = pos_before + 1; - while bytes[current] == 1 { - current += 2; - } - let to_skip = current - pos_before + 1; - deserializer.advance(to_skip); - return Ok(to_skip); - } - ConcreteDataType::String(_) => { - let pos_before = deserializer.position(); - deserializer.advance(1); - deserializer - .skip_bytes() - .context(error::DeserializeFieldSnafu)?; - return Ok(deserializer.position() - pos_before); - } - ConcreteDataType::Date(_) => 5, - ConcreteDataType::Timestamp(_) => 9, // We treat timestamp as Option - ConcreteDataType::Time(_) => 10, // i64 and 1 byte time unit - ConcreteDataType::Duration(_) => 10, + ConcreteDataType::Int32(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::Float32(_) + | ConcreteDataType::Date(_) => 5, + ConcreteDataType::Int64(_) + | ConcreteDataType::UInt64(_) + | ConcreteDataType::Float64(_) + | ConcreteDataType::Timestamp(_) => 9, + ConcreteDataType::Time(_) | ConcreteDataType::Duration(_) => 10, ConcreteDataType::Interval(IntervalType::YearMonth(_)) => 5, ConcreteDataType::Interval(IntervalType::DayTime(_)) => 9, ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => 17, ConcreteDataType::Decimal128(_) => 19, - ConcreteDataType::Null(_) - | ConcreteDataType::List(_) - | ConcreteDataType::Struct(_) - | ConcreteDataType::Dictionary(_) => 0, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) => { + // A serde sequence: each byte is preceded by 1, terminated by 0. + let mut pos = 1; + while bytes.get(pos) == Some(&1) { + pos += 2; + } + if bytes.get(pos) != Some(&0) { + return Err(invalid()); + } + pos + 1 + } + ConcreteDataType::String(_) => { + match bytes.get(1) { + Some(0) => return Ok(2), + Some(1) => {} + _ => return Err(invalid()), + } + let mut pos = 2; + loop { + match bytes.get(pos + 8) { + Some(1..=8) => break pos + 9, + Some(9) => pos += 9, + _ => return Err(invalid()), + } + } + } + data_type => { + return error::NotSupportedFieldSnafu { + data_type: data_type.clone(), + } + .fail(); + } }; - deserializer.advance(to_skip); - Ok(to_skip) + snafu::ensure!( + bytes.len() >= len, + error::InvalidDensePrimaryKeySnafu { + reason: "truncated field", + } + ); + Ok(len) } } @@ -411,6 +430,35 @@ impl DensePrimaryKeyCodec { Ok(values) } + /// Counts complete fields in a key encoded with a prefix of this codec's schema. + /// + /// Dense keys contain neither column ids nor types: callers must ensure that + /// existing fields have the same order and types. Only EOF between fields is + /// accepted; truncated fields and bytes beyond the schema return an error. + /// Field boundaries are checked in all builds; full value validation is only + /// enabled in debug builds and this crate's unit tests. + pub fn decode_prefix_len(&self, bytes: &[u8]) -> Result { + let mut deserializer = Deserializer::new(bytes); + for (index, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() { + if !deserializer.has_remaining() { + return Ok(index); + } + let start = deserializer.position(); + let len = field.encoded_field_len(&bytes[start..])?; + // Production range mapping only needs boundaries, not allocated field values. + #[cfg(any(debug_assertions, test))] + field.deserialize(&mut Deserializer::new(&bytes[start..start + len]))?; + deserializer.advance(len); + } + snafu::ensure!( + !deserializer.has_remaining(), + error::InvalidDensePrimaryKeySnafu { + reason: "key contains bytes beyond the primary key schema", + } + ); + Ok(self.num_fields()) + } + /// Decode primary key values from bytes without column id. pub fn decode_dense_without_column_id(&self, bytes: &[u8]) -> Result> { let mut deserializer = Deserializer::new(bytes); @@ -594,6 +642,25 @@ mod tests { let value_ref = row.iter().map(|v| v.as_value_ref()).collect::>(); let result = encoder.encode(value_ref.iter().cloned()).unwrap(); + let boundaries: Vec<_> = (0..=row.len()) + .map(|count| { + encoder + .encode(value_ref[..count].iter().cloned()) + .unwrap() + .len() + }) + .collect(); + for end in 0..=result.len() { + match boundaries.iter().position(|boundary| *boundary == end) { + Some(count) => { + assert_eq!(count, encoder.decode_prefix_len(&result[..end]).unwrap()) + } + None => assert!( + encoder.decode_prefix_len(&result[..end]).is_err(), + "truncated {data_types:?} at {end}" + ), + } + } let decoded = encoder.decode(&result).unwrap().into_dense(); assert_eq!(decoded, row); let mut decoded = Vec::new(); @@ -610,6 +677,61 @@ mod tests { } } + #[test] + fn test_decode_prefix_len_accepts_only_complete_fields() { + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ConcreteDataType::string_datatype())), + (1, SortField::new(ConcreteDataType::int64_datatype())), + (2, SortField::new(ConcreteDataType::binary_datatype())), + (3, SortField::new(ConcreteDataType::string_datatype())), + ]); + let values = [ + Value::from("abcdefghijk"), + Value::Int64(42), + Value::Binary(vec![0, 255, 1].into()), + Value::Null, + ]; + let mut boundaries = vec![0]; + let mut encoded = Vec::new(); + for count in 1..=values.len() { + encoded.clear(); + codec + .encode_dense( + values[..count].iter().map(Value::as_value_ref), + &mut encoded, + ) + .unwrap(); + boundaries.push(encoded.len()); + assert_eq!(count, codec.decode_prefix_len(&encoded).unwrap()); + } + for end in 0..=encoded.len() { + match boundaries.iter().position(|boundary| *boundary == end) { + Some(count) => assert_eq!(count, codec.decode_prefix_len(&encoded[..end]).unwrap()), + None => assert!( + codec.decode_prefix_len(&encoded[..end]).is_err(), + "truncated at {end}" + ), + } + } + encoded.push(0); + assert!(codec.decode_prefix_len(&encoded).is_err()); + assert!(codec.decode_prefix_len(&[2]).is_err()); + let empty = DensePrimaryKeyCodec::with_fields(vec![]); + assert_eq!(0, empty.decode_prefix_len(&[]).unwrap()); + assert!(empty.decode_prefix_len(&[0]).is_err()); + } + + #[test] + fn test_prefix_len_validates_values_in_unit_tests() { + let codec = DensePrimaryKeyCodec::with_fields(vec![( + 0, + SortField::new(ConcreteDataType::boolean_datatype()), + )]); + // The field is complete, but 2 is not a boolean. Even release unit tests + // retain value validation through cfg(test). + assert!(codec.decode_prefix_len(&[1, 2]).is_err()); + } + #[test] fn test_memcmp() { let encoder = DensePrimaryKeyCodec::with_fields(vec![ diff --git a/src/mito2/benches/bench_compaction_picker.rs b/src/mito2/benches/bench_compaction_picker.rs index 439b9e6f226..99f02831536 100644 --- a/src/mito2/benches/bench_compaction_picker.rs +++ b/src/mito2/benches/bench_compaction_picker.rs @@ -78,7 +78,7 @@ fn bench_find_sorted_runs(c: &mut Criterion) { group.bench_function(format!("{}_new", case_name), |b| { let mut files = generate_same_timestamp_files(total_files, files_per_timestamp); b.iter(|| { - find_sorted_runs(black_box(&mut files)); + find_sorted_runs(black_box(&mut files), Ranged::overlap); }); }); @@ -121,7 +121,12 @@ fn bench_find_overlapping_items(c: &mut Criterion) { let mut r2 = SortedRun::from(files2); b.iter(|| { let mut result = vec![]; - find_overlapping_items(black_box(&mut r1), black_box(&mut r2), &mut result); + find_overlapping_items( + black_box(&mut r1), + black_box(&mut r2), + &mut result, + Ranged::overlap_inclusive, + ); }); }); } diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 0a7bb61813b..c7d55fb8a28 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -208,7 +208,7 @@ pub async fn open_compaction_region( }; let current_version = { - let mut ssts = SstVersion::new(); + let mut ssts = SstVersion::new(region_metadata.clone()); ssts.add_files(file_purger.clone(), manifest.files.values().cloned()); CompactionVersion { metadata: region_metadata.clone(), @@ -1170,9 +1170,9 @@ mod tests { access_layer: env.access_layer.clone(), manifest_ctx, current_version: CompactionVersion { - metadata, + metadata: metadata.clone(), options: RegionOptions::default(), - ssts: Arc::new(SstVersion::new()), + ssts: Arc::new(SstVersion::new(metadata)), memtable_min_sequence: None, compaction_time_window: None, }, diff --git a/src/mito2/src/compaction/last_non_null.rs b/src/mito2/src/compaction/last_non_null.rs index 60f9bc254af..afcdc87e90f 100644 --- a/src/mito2/src/compaction/last_non_null.rs +++ b/src/mito2/src/compaction/last_non_null.rs @@ -15,7 +15,7 @@ use std::collections::{HashMap, HashSet}; use snafu::ResultExt; -use store_api::metadata::RegionMetadata; +use store_api::metadata::RegionMetadataRef; use store_api::storage::SequenceNumber; use crate::compaction::CompactionOutput; @@ -122,7 +122,7 @@ pub(super) fn inputs_precede_memtables( fn build_closed_outputs( seeds: Vec, files: Vec, - metadata: &RegionMetadata, + metadata: &RegionMetadataRef, ) -> Vec { // Maps every seed input to its seed, so a closure that reaches that input // can absorb the rest of the seed instead of letting it schedule separately. @@ -406,12 +406,12 @@ mod tests { #[case("dense_with_external", 1024)] #[case("time_dense_pk_disjoint", 1)] fn test_large_snapshot_closure(#[case] shape: &str, #[case] expected_count: usize) { - use bytes::Bytes; + use crate::compaction::test_util::{ + new_file_handle_with_size_sequence_and_primary_key_range, pk_range, + primary_key_metadata_for_test, + }; - use crate::compaction::test_util::new_file_handle_with_size_sequence_and_primary_key_range; - - let mut metadata = (*metadata_for_test()).clone(); - metadata.region_id = 0.into(); + let metadata = primary_key_metadata_for_test(); let files = (0..2048) .map(|i| { let (start, end, pk) = match shape { @@ -419,8 +419,8 @@ mod tests { "dense_with_external" if i < 1024 => (0, 1, None), "dense_with_external" => (i, i, None), "time_dense_pk_disjoint" => { - let pk = Bytes::copy_from_slice(&(i as u32).to_be_bytes()); - (0, 1, Some((pk.clone(), pk))) + let key = i.to_string(); + (0, 1, pk_range(key.as_bytes(), key.as_bytes())) } _ => unreachable!(), }; diff --git a/src/mito2/src/compaction/overlap.rs b/src/mito2/src/compaction/overlap.rs index 38cb7ad8a4a..f06c9058455 100644 --- a/src/mito2/src/compaction/overlap.rs +++ b/src/mito2/src/compaction/overlap.rs @@ -16,10 +16,11 @@ use std::collections::HashMap; use std::ops::Range; use common_time::Timestamp; -use store_api::metadata::RegionMetadata; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use crate::compaction::run::primary_key_ranges_overlap; use crate::sst::file::{FileHandle, RegionFileId}; +use crate::sst::primary_key::PrimaryKeyRangeMapper; /// Snapshot-local overlap candidates, ordered by start time. Each subtree stores /// its maximum active end. Removing visited files prunes dense internal overlaps. @@ -28,6 +29,7 @@ use crate::sst::file::{FileHandle, RegionFileId}; pub(super) struct FileOverlapIndex<'a> { /// Current region metadata used to decide whether PK bounds can safely exclude overlaps. metadata: &'a RegionMetadata, + primary_key_mapper: PrimaryKeyRangeMapper, /// Candidates in original snapshot order, retained after removal so their /// indices remain stable and drained matches can preserve merge input order. files: Vec, @@ -54,7 +56,7 @@ struct OverlapQuery<'a> { } impl<'a> FileOverlapIndex<'a> { - pub(super) fn new(files: Vec, metadata: &'a RegionMetadata) -> Self { + pub(super) fn new(files: Vec, metadata: &'a RegionMetadataRef) -> Self { let mut by_start: Vec<_> = (0..files.len()).collect(); by_start.sort_unstable_by_key(|&i| (files[i].time_range().0, i)); // A single empty leaf also handles an empty snapshot. @@ -70,6 +72,7 @@ impl<'a> FileOverlapIndex<'a> { } Self { metadata, + primary_key_mapper: PrimaryKeyRangeMapper::new(metadata.clone()), files, by_start, positions, @@ -127,7 +130,13 @@ impl<'a> FileOverlapIndex<'a> { } if leaves.len() == 1 { let i = self.by_start[leaves.start]; - return files_may_overlap(query.input, &self.files[i], self.metadata).then_some(i); + return files_may_overlap( + query.input, + &self.files[i], + self.metadata, + &self.primary_key_mapper, + ) + .then_some(i); } let mid = leaves.start + leaves.len() / 2; self.find_in_subtree(2 * node, leaves.start..mid, query) @@ -137,7 +146,12 @@ impl<'a> FileOverlapIndex<'a> { /// SST bounds are inclusive. Missing statistics, foreign encodings and schema /// evolution must not exclude a possible logical-key dependency. -fn files_may_overlap(lhs: &FileHandle, rhs: &FileHandle, metadata: &RegionMetadata) -> bool { +fn files_may_overlap( + lhs: &FileHandle, + rhs: &FileHandle, + metadata: &RegionMetadata, + mapper: &PrimaryKeyRangeMapper, +) -> bool { let (lhs_start, lhs_end) = lhs.time_range(); let (rhs_start, rhs_end) = rhs.time_range(); if lhs_start.max(rhs_start) > lhs_end.min(rhs_end) { @@ -155,7 +169,7 @@ fn files_may_overlap(lhs: &FileHandle, rhs: &FileHandle, metadata: &RegionMetada { return true; } - match (lhs.primary_key_range(), rhs.primary_key_range()) { + match (lhs.primary_key_range(mapper), rhs.primary_key_range(mapper)) { (Some(lhs), Some(rhs)) if lhs.0 <= lhs.1 && rhs.0 <= rhs.1 => { primary_key_ranges_overlap(&lhs, &rhs) } @@ -166,12 +180,13 @@ fn files_may_overlap(lhs: &FileHandle, rhs: &FileHandle, metadata: &RegionMetada #[cfg(test)] mod tests { use std::collections::HashSet; + use std::sync::Arc; - use bytes::Bytes; use rand::{Rng, SeedableRng}; use store_api::storage::FileId; use super::*; + use crate::compaction::test_util::{pk_range, primary_key_metadata_for_test}; use crate::sst::file::FileMeta; use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::new_noop_file_purger; @@ -187,12 +202,7 @@ mod tests { ..Default::default() }, new_noop_file_purger(), - pk.map(|(start, end)| { - ( - Bytes::copy_from_slice(start.as_bytes()), - Bytes::copy_from_slice(end.as_bytes()), - ) - }), + pk.and_then(|(start, end)| pk_range(start.as_bytes(), end.as_bytes())), ) } @@ -212,9 +222,11 @@ mod tests { #[case] foreign: bool, #[case] expected: bool, ) { - let mut metadata = (*metadata_for_test()).clone(); + let mut metadata = (*primary_key_metadata_for_test()).clone(); metadata.region_id = 0.into(); metadata.schema_version = schema_version; + let metadata = Arc::new(metadata); + let ranges = PrimaryKeyRangeMapper::new(metadata.clone()); let lhs = file(0, 10, Some(("a", "b"))); let mut rhs = file(start, start + 10, pk); if foreign { @@ -223,11 +235,11 @@ mod tests { rhs = FileHandle::new_with_primary_key_range( meta, new_noop_file_purger(), - rhs.primary_key_range(), + rhs.raw_primary_key_range(), ); } - assert_eq!(expected, files_may_overlap(&lhs, &rhs, &metadata)); - assert_eq!(expected, files_may_overlap(&rhs, &lhs, &metadata)); + assert_eq!(expected, files_may_overlap(&lhs, &rhs, &metadata, &ranges)); + assert_eq!(expected, files_may_overlap(&rhs, &lhs, &metadata, &ranges)); } #[test] @@ -247,8 +259,8 @@ mod tests { #[test] fn test_index_matches_linear_scan_after_removals() { - let mut metadata = (*metadata_for_test()).clone(); - metadata.region_id = 0.into(); + let metadata = primary_key_metadata_for_test(); + let ranges = PrimaryKeyRangeMapper::new(metadata.clone()); let mut rng = rand::rngs::StdRng::seed_from_u64(9146); for count in [0, 1, 7, 32, 127] { let files: Vec<_> = (0..count) @@ -276,7 +288,8 @@ mod tests { let expected: Vec<_> = files .iter() .filter(|f| { - !removed.contains(&f.file_id()) && files_may_overlap(&query, f, &metadata) + !removed.contains(&f.file_id()) + && files_may_overlap(&query, f, &metadata, &ranges) }) .map(FileHandle::file_id) .collect(); diff --git a/src/mito2/src/compaction/picker.rs b/src/mito2/src/compaction/picker.rs index de2e9d9d310..ec9fa91f4c8 100644 --- a/src/mito2/src/compaction/picker.rs +++ b/src/mito2/src/compaction/picker.rs @@ -91,7 +91,8 @@ impl From<&PickerOutput> for SerializedPickerOutput { } impl PickerOutput { - /// Converts a [SerializedPickerOutput] to a [PickerOutput]. + /// Converts a [SerializedPickerOutput] to a [PickerOutput]. File statistics + /// retain their original encoding; comparisons supply their own schema context. pub fn from_serialized( input: SerializedPickerOutput, file_purger: Arc, diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 5698a6ce667..2ad9f2f1eba 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -23,6 +23,7 @@ use common_base::BitVec; use common_time::Timestamp; use crate::sst::file::FileHandle; +use crate::sst::primary_key::PrimaryKeyRangeMapper; /// Trait for any items with specific range (both boundaries are inclusive). pub trait Ranged { @@ -64,10 +65,12 @@ pub(crate) fn merge_primary_key_ranges( } } +/// Uses the caller's inclusive overlap test after pruning disjoint time ranges. pub fn find_overlapping_items( l: &mut SortedRun, r: &mut SortedRun, result: &mut Vec, + overlaps: impl Fn(&T, &T) -> bool, ) { if l.items.is_empty() || r.items.is_empty() { return; @@ -114,7 +117,7 @@ pub fn find_overlapping_items( } // We have an overlap (inclusive: touching boundaries count) - if lhs.overlap_inclusive(&r.items[j]) { + if overlaps(lhs, &r.items[j]) { if !selected[lhs_idx] { result.push(lhs.clone()); selected.set(lhs_idx, true); @@ -147,37 +150,41 @@ pub trait Item: Ranged + Clone { fn size(&self) -> usize; } +// Physical handles only supply time ranges. PK-aware comparisons need a target schema. impl Ranged for FileHandle { type BoundType = Timestamp; fn range(&self) -> (Self::BoundType, Self::BoundType) { self.time_range() } +} - fn overlap(&self, other: &Self) -> bool { - let (lhs_start, lhs_end) = self.range(); - let (rhs_start, rhs_end) = other.range(); - if lhs_start.max(rhs_start) >= lhs_end.min(rhs_end) { - return false; - } +/// Tests exclusive time overlap and inclusive PK overlap in one pinned schema. +pub(crate) fn files_overlap( + lhs: &FileHandle, + rhs: &FileHandle, + mapper: &PrimaryKeyRangeMapper, +) -> bool { + lhs.overlap(rhs) && file_primary_keys_overlap(lhs, rhs, mapper) +} - match (&self.primary_key_range(), &other.primary_key_range()) { - (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), - _ => true, - } - } +/// Includes touching time and PK boundaries when checking file dependencies. +pub(crate) fn files_overlap_inclusive( + lhs: &FileHandle, + rhs: &FileHandle, + mapper: &PrimaryKeyRangeMapper, +) -> bool { + lhs.overlap_inclusive(rhs) && file_primary_keys_overlap(lhs, rhs, mapper) +} - fn overlap_inclusive(&self, other: &Self) -> bool { - let (lhs_start, lhs_end) = self.range(); - let (rhs_start, rhs_end) = other.range(); - if lhs_start.max(rhs_start) > lhs_end.min(rhs_end) { - return false; - } - - match (&self.primary_key_range(), &other.primary_key_range()) { - (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), - _ => true, - } +fn file_primary_keys_overlap( + lhs: &FileHandle, + rhs: &FileHandle, + mapper: &PrimaryKeyRangeMapper, +) -> bool { + match (lhs.primary_key_range(mapper), rhs.primary_key_range(mapper)) { + (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(&lhs, &rhs), + _ => true, } } @@ -251,8 +258,8 @@ where } } -/// Finds sorted runs in given items. -pub fn find_sorted_runs(items: &mut [T]) -> Vec> +/// Finds sorted runs using the caller's overlap test within each active time range. +pub fn find_sorted_runs(items: &mut [T], overlaps: impl Fn(&T, &T) -> bool) -> Vec> where T: Item, { @@ -296,7 +303,7 @@ where let mut overlaps_any = false; for idx in &active_run_item_indices { let run_item = ¤t_run.items[*idx]; - if run_item.overlap(item) { + if overlaps(run_item, item) { overlaps_any = true; break; } @@ -490,11 +497,13 @@ where #[cfg(test)] mod tests { - use bytes::Bytes; use store_api::storage::FileId; use super::*; - use crate::compaction::test_util::new_file_handle_with_size_sequence_and_primary_key_range; + use crate::compaction::test_util::{ + new_file_handle_with_size_sequence_and_primary_key_range, pk_range, + primary_key_mapper_for_test, + }; #[derive(Clone, Debug, PartialEq)] struct MockFile { @@ -528,10 +537,6 @@ mod tests { .collect() } - fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> { - Some((Bytes::from_static(min), Bytes::from_static(max))) - } - fn check_sorted_runs( ranges: &[(i64, i64)], expected_runs: &[Vec<(i64, i64)>], @@ -539,7 +544,7 @@ mod tests { let mut files = build_items(ranges); let mut files_clone = files.clone(); - let runs = find_sorted_runs(&mut files); + let runs = find_sorted_runs(&mut files, Ranged::overlap); let result_file_ranges: Vec> = runs .iter() @@ -574,7 +579,7 @@ mod tests { let mut files = build_items(ranges); let mut files_for_original = files.clone(); - let runs = find_sorted_runs(&mut files); + let runs = find_sorted_runs(&mut files, Ranged::overlap); let original_runs = find_sorted_runs_original(&mut files_for_original); assert_eq!(sorted_run_ranges(&original_runs), sorted_run_ranges(&runs)); @@ -661,6 +666,7 @@ mod tests { &mut SortedRun::from(Vec::::new()), &mut SortedRun::from(Vec::::new()), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result, Vec::::new()); @@ -669,6 +675,7 @@ mod tests { &mut SortedRun::from(files1.clone()), &mut SortedRun::from(Vec::::new()), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result, Vec::::new()); @@ -676,6 +683,7 @@ mod tests { &mut SortedRun::from(Vec::::new()), &mut SortedRun::from(files1.clone()), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result, Vec::::new()); @@ -686,6 +694,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result, Vec::::new()); @@ -696,6 +705,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 2); assert_eq!(result[0].range(), (1, 5)); @@ -708,6 +718,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 6); @@ -718,6 +729,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 2); // Should overlap since ranges are inclusive @@ -728,6 +740,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 2); @@ -738,6 +751,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 2); @@ -748,6 +762,7 @@ mod tests { &mut SortedRun::from(files1), &mut SortedRun::from(files2), &mut result, + Ranged::overlap_inclusive, ); assert_eq!(result.len(), 4); // Should find both overlaps } @@ -773,7 +788,7 @@ mod tests { pk_range(b"x", b"z"), ); - assert!(!lhs.overlap(&rhs)); + assert!(!files_overlap(&lhs, &rhs, &primary_key_mapper_for_test())); } #[test] @@ -799,7 +814,8 @@ mod tests { ), ]; - let runs = find_sorted_runs(&mut files); + let ranges = primary_key_mapper_for_test(); + let runs = find_sorted_runs(&mut files, |lhs, rhs| files_overlap(lhs, rhs, &ranges)); assert_eq!(1, runs.len()); assert_eq!(2, runs[0].items().len()); @@ -837,7 +853,8 @@ mod tests { ), ]; - let runs = find_sorted_runs(&mut files); + let ranges = primary_key_mapper_for_test(); + let runs = find_sorted_runs(&mut files, |lhs, rhs| files_overlap(lhs, rhs, &ranges)); assert_eq!(2, runs.len()); assert_eq!(2, runs[0].items().len()); @@ -870,7 +887,10 @@ mod tests { ]); let mut result = Vec::new(); - find_overlapping_items(&mut left, &mut right, &mut result); + let ranges = primary_key_mapper_for_test(); + find_overlapping_items(&mut left, &mut right, &mut result, |lhs, rhs| { + files_overlap_inclusive(lhs, rhs, &ranges) + }); assert!(result.is_empty()); } @@ -896,6 +916,6 @@ mod tests { pk_range(b"a", b"f"), ); - assert!(!lhs.overlap(&rhs)); + assert!(!files_overlap(&lhs, &rhs, &primary_key_mapper_for_test())); } } diff --git a/src/mito2/src/compaction/test_util.rs b/src/mito2/src/compaction/test_util.rs index 453af8a8a68..1744b37c95b 100644 --- a/src/mito2/src/compaction/test_util.rs +++ b/src/mito2/src/compaction/test_util.rs @@ -19,6 +19,7 @@ use std::time::Duration; use bytes::Bytes; use common_base::Plugins; use common_time::Timestamp; +use store_api::metadata::RegionMetadataRef; use store_api::storage::FileId; use crate::cache::CacheManager; @@ -26,11 +27,30 @@ use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; use crate::config::MitoConfig; use crate::region::options::RegionOptions; use crate::sst::file::{FileHandle, FileMeta, Level}; +use crate::sst::primary_key::PrimaryKeyRangeMapper; use crate::sst::version::SstVersion; use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::new_noop_file_purger; use crate::test_util::scheduler_util::SchedulerEnv; +pub(crate) fn primary_key_metadata_for_test() -> RegionMetadataRef { + let mut metadata = crate::test_util::memtable_util::metadata_with_primary_key(vec![0], false); + metadata.region_id = 0.into(); + Arc::new(metadata) +} + +pub(crate) fn primary_key_mapper_for_test() -> PrimaryKeyRangeMapper { + PrimaryKeyRangeMapper::new(primary_key_metadata_for_test()) +} + +/// Encodes the single string tag used by compaction range fixtures. +pub(crate) fn pk_range(min: &[u8], max: &[u8]) -> Option<(Bytes, Bytes)> { + let encode = |key| { + crate::test_util::sst_util::new_primary_key(&[std::str::from_utf8(key).unwrap()]).into() + }; + Some((encode(min), encode(max))) +} + /// Test util to create file handles. pub fn new_file_handle( file_id: FileId, @@ -144,9 +164,12 @@ pub(crate) async fn compaction_region_with_ssts( ttl: Duration, ) -> CompactionRegion { let env = SchedulerEnv::new().await; - let metadata = metadata_for_test(); + let mut metadata = (*metadata_for_test()).clone(); + // Match the table used by new_file_handle* and default FileMeta fixtures. + metadata.region_id = 0.into(); + let metadata = Arc::new(metadata); let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - let mut ssts = SstVersion::new(); + let mut ssts = SstVersion::new(metadata.clone()); ssts.add_files( Arc::new(crate::sst::file_purger::NoopFilePurger), files.into_iter(), diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index b2e3274b062..dd407b30382 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -31,11 +31,12 @@ use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::compaction::run::{ - Ranged, SortedRun, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, - primary_key_ranges_overlap, + Ranged, SortedRun, files_overlap, files_overlap_inclusive, find_sorted_runs, + find_sorted_runs_by_time_range, merge_primary_key_ranges, primary_key_ranges_overlap, }; use crate::error::{JoinSnafu, Result}; use crate::sst::file::{FileHandle, Level, overlaps}; +use crate::sst::primary_key::PrimaryKeyRangeMapper; use crate::sst::version::LevelMeta; const LEVEL_COMPACTED: Level = 1; @@ -69,12 +70,14 @@ struct WindowPickContext<'a> { files: &'a Window, windows: &'a BTreeMap, phase: PickPhase, + primary_key_mapper: &'a PrimaryKeyRangeMapper, } struct WindowOutputContext { active_window: Option, time_window_size: Option, max_outputs: Option, + primary_key_mapper: Arc, } /// A mixed L0/L1 compaction may rewrite at most this many L1 rows per L0 row. @@ -135,6 +138,7 @@ impl TwcsPicker { active_window, time_window_size, max_outputs, + primary_key_mapper, } = context; let mut output = vec![]; let windows = time_windows @@ -175,6 +179,7 @@ impl TwcsPicker { let picker = self.clone(); let time_windows = time_windows.clone(); let window = *window; + let primary_key_mapper = primary_key_mapper.clone(); handles.push(common_runtime::spawn_blocking_compact(move || { time_windows.get(&window).map(|files| { ( @@ -186,6 +191,7 @@ impl TwcsPicker { files, windows: &time_windows, phase, + primary_key_mapper: &primary_key_mapper, }, ), ) @@ -238,6 +244,7 @@ impl TwcsPicker { files, windows, phase, + primary_key_mapper, } = context; let is_active_window = active_window == Some(files.time_window); let window = &files.time_window; @@ -275,13 +282,21 @@ impl TwcsPicker { } let (inputs, found_runs) = if is_active_window { match phase { - PickPhase::HasL0 if num_l0_files >= self.trigger_file_num => { - pick_candidate_files(l0_files, self.max_output_file_size, pick_count_first) - } + PickPhase::HasL0 if num_l0_files >= self.trigger_file_num => pick_candidate_files( + l0_files, + self.max_output_file_size, + pick_count_first, + primary_key_mapper, + ), PickPhase::L1FileReduction | PickPhase::L1OverlapOnly if num_l1_files >= self.active_window_l1_merge_trigger => { - pick_l1_candidate_files(l1_files, self.max_output_file_size, phase) + pick_l1_candidate_files( + l1_files, + self.max_output_file_size, + phase, + primary_key_mapper, + ) } _ => (vec![], 0), } @@ -293,6 +308,7 @@ impl TwcsPicker { l0_file_num: self.inactive_window_trigger_file_num, l1_file_num: self.inactive_window_l1_merge_trigger, phase, + primary_key_mapper, }, self.max_output_file_size, ) @@ -303,7 +319,7 @@ impl TwcsPicker { let filter_deleted = !self.append_mode && !window_has_overlap(files, windows) - && !selected_overlaps_unselected(&inputs, files); + && !selected_overlaps_unselected(&inputs, files, primary_key_mapper); if inputs.len() > 1 { // If we have more than one file to compact. @@ -336,72 +352,106 @@ impl TwcsPicker { /// The rewrite is bounded by the L0 bytes and leaves large compacted files /// untouched. If nothing qualifies, the window is left as-is. #[derive(Debug, Clone, Copy)] -struct InactiveWindowPick { +struct InactiveWindowPick<'a> { l0_file_num: usize, l1_file_num: usize, phase: PickPhase, + primary_key_mapper: &'a PrimaryKeyRangeMapper, } fn pick_inactive_window_files( l0_files: Vec, l1_files: Vec, - pick: InactiveWindowPick, + pick: InactiveWindowPick<'_>, max_output_file_size: Option, ) -> (Vec, usize) { if !matches!(pick.phase, PickPhase::HasL0) { return if l1_files.len() >= pick.l1_file_num { - pick_l1_candidate_files(l1_files, max_output_file_size, pick.phase) + pick_l1_candidate_files( + l1_files, + max_output_file_size, + pick.phase, + pick.primary_key_mapper, + ) } else { (vec![], 0) }; } if l0_files.len() >= pick.l0_file_num { - let pick = pick_candidate_files(l0_files.clone(), max_output_file_size, pick_count_first); - if !pick.0.is_empty() { - return pick; + let candidate = pick_candidate_files( + l0_files.clone(), + max_output_file_size, + pick_count_first, + pick.primary_key_mapper, + ); + if !candidate.0.is_empty() { + return candidate; } } - let pick = pick_candidate_files(l0_files.clone(), max_output_file_size, pick_count_first); - if !pick.0.is_empty() { - return pick; + let candidate = pick_candidate_files( + l0_files.clone(), + max_output_file_size, + pick_count_first, + pick.primary_key_mapper, + ); + if !candidate.0.is_empty() { + return candidate; } let mut all_files = l0_files.clone(); all_files.extend(l1_files); - let pick = pick_candidate_files(all_files, max_output_file_size, pick_mixed_within_budget); - if !pick.0.is_empty() { - return pick; + let candidate = pick_candidate_files( + all_files, + max_output_file_size, + pick_mixed_within_budget, + pick.primary_key_mapper, + ); + if !candidate.0.is_empty() { + return candidate; } - pick_candidate_files(l0_files, max_output_file_size, pick_unbalanced_count_first) + pick_candidate_files( + l0_files, + max_output_file_size, + pick_unbalanced_count_first, + pick.primary_key_mapper, + ) } fn pick_l1_candidate_files( l1_files: Vec, max_output_file_size: Option, phase: PickPhase, + mapper: &PrimaryKeyRangeMapper, ) -> (Vec, usize) { let picker = match phase { PickPhase::L1FileReduction => pick_l1_file_reduction, PickPhase::L1OverlapOnly => pick_l1_overlap_only, PickPhase::HasL0 => return (vec![], 0), }; - pick_candidate_files(l1_files, max_output_file_size, picker) + pick_candidate_files(l1_files, max_output_file_size, picker, mapper) } +type CandidatePicker = + fn(Vec>, Option, &PrimaryKeyRangeMapper) -> Vec; + fn pick_candidate_files( mut files: Vec, max_output_file_size: Option, - picker: fn(Vec>, Option) -> Vec, + picker: CandidatePicker, + mapper: &PrimaryKeyRangeMapper, ) -> (Vec, usize) { let sorted_runs = if files.len() < 1024 { - find_sorted_runs(&mut files) + find_sorted_runs(&mut files, |lhs, rhs| files_overlap(lhs, rhs, mapper)) } else { find_sorted_runs_by_time_range(&mut files) }; let found_runs = sorted_runs.len(); - (picker(sorted_runs, max_output_file_size), found_runs) + ( + picker(sorted_runs, max_output_file_size, mapper), + found_runs, + ) } #[derive(Debug)] @@ -443,6 +493,7 @@ impl Candidate { file: &OrderedFile, preceding: &[OrderedFile], participations: &mut Vec, + mapper: &PrimaryKeyRangeMapper, ) { self.num_files += 1; let file_size = file.file.size() as usize; @@ -452,7 +503,8 @@ impl Candidate { let mut participates = false; for (offset, other) in preceding.iter().enumerate() { - if file.run_id != other.run_id && file.file.overlap_inclusive(other.file) { + if file.run_id != other.run_id && files_overlap_inclusive(file.file, other.file, mapper) + { if !participations[offset] { participations[offset] = true; self.overlap_participants += 1; @@ -582,15 +634,22 @@ impl CandidateScore { fn pick_count_first( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, is_balanced_candidate) + pick_count_first_where( + sorted_runs, + max_output_file_size, + mapper, + is_balanced_candidate, + ) } fn pick_l1_file_reduction( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, |candidate| { + pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { is_balanced_candidate(candidate) && candidate.file_reduction(max_output_file_size) > 0 }) } @@ -598,8 +657,9 @@ fn pick_l1_file_reduction( fn pick_l1_overlap_only( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, |candidate| { + pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { is_balanced_candidate(candidate) && candidate.file_reduction(max_output_file_size) == 0 && candidate.overlap_participants > 0 @@ -610,8 +670,9 @@ fn pick_l1_overlap_only( fn pick_mixed_count_first( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, |candidate| { + pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { candidate.has_mixed_levels() && is_balanced_candidate(candidate) }) } @@ -621,8 +682,9 @@ fn pick_mixed_count_first( fn pick_mixed_within_budget( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, |candidate| { + pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { candidate.has_mixed_levels() && candidate.within_rewrite_budget(max_output_file_size) }) } @@ -633,8 +695,9 @@ fn pick_mixed_within_budget( fn pick_unbalanced_count_first( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, ) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, |_| true) + pick_count_first_where(sorted_runs, max_output_file_size, mapper, |_| true) } fn is_balanced_candidate(candidate: &Candidate) -> bool { @@ -644,6 +707,7 @@ fn is_balanced_candidate(candidate: &Candidate) -> bool { fn pick_count_first_where( sorted_runs: Vec>, max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, is_eligible: impl Fn(&Candidate) -> bool, ) -> Vec { let files = ordered_files(&sorted_runs); @@ -654,7 +718,12 @@ fn pick_count_first_where( let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len()); let mut participations: Vec = Vec::with_capacity(right_bound - left); for right in left..right_bound { - candidate.absorb(&files[right], &files[left..right], &mut participations); + candidate.absorb( + &files[right], + &files[left..right], + &mut participations, + mapper, + ); if candidate.num_files < 2 || !is_eligible(&candidate) || !candidate.makes_progress(max_output_file_size) @@ -707,7 +776,11 @@ fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> files } -fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> bool { +fn selected_overlaps_unselected( + selected: &[FileHandle], + window: &Window, + mapper: &PrimaryKeyRangeMapper, +) -> bool { // The overall time span of the selection: a file outside it cannot overlap any // selected file (ranges are inclusive), so it needs no precise overlap check. let Some((span_start, span_end)) = selected @@ -731,7 +804,7 @@ fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> boo .any(|unselected| { selected .iter() - .any(|selected| selected.overlap_inclusive(unselected)) + .any(|selected| files_overlap_inclusive(selected, unselected, mapper)) }) } @@ -798,6 +871,8 @@ impl TwcsPicker { let region_id = compaction_region.region_id; let picker = self.clone(); let compaction_region = compaction_region.clone(); + let primary_key_mapper = compaction_region.current_version.ssts.primary_key_mapper(); + let window_mapper = primary_key_mapper.clone(); let (expired_ssts, time_window_size, active_window, windows) = common_runtime::spawn_blocking_compact(move || { let levels = compaction_region.current_version.ssts.levels(); @@ -835,6 +910,7 @@ impl TwcsPicker { .flat_map(LevelMeta::files) .filter(|file| !expired_file_ids.contains(&file.file_id())), time_window_size, + &window_mapper, ); // Compute activity from the candidate files so expired or // compacting files cannot identify a window absent from `windows`. @@ -865,6 +941,7 @@ impl TwcsPicker { active_window, time_window_size: Some(time_window_size), max_outputs, + primary_key_mapper, }, ) .await?; @@ -898,9 +975,9 @@ struct Window { impl Window { /// Creates a new [Window] with given file. - fn new_with_file(file: FileHandle) -> Self { + fn new_with_file(file: FileHandle, mapper: &PrimaryKeyRangeMapper) -> Self { let (start, end) = file.time_range(); - let primary_key_range = file.primary_key_range(); + let primary_key_range = file.primary_key_range(mapper); Self { start, end, @@ -916,12 +993,14 @@ impl Window { } /// Adds a new file to window and updates time range. - fn add_file(&mut self, file: FileHandle) { + fn add_file(&mut self, file: FileHandle, mapper: &PrimaryKeyRangeMapper) { let (start, end) = file.time_range(); self.start = self.start.min(start); self.end = self.end.max(end); - self.primary_key_range = - merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); + self.primary_key_range = merge_primary_key_ranges( + self.primary_key_range.take(), + file.primary_key_range(mapper), + ); self.files.push(file); } @@ -934,6 +1013,7 @@ impl Window { fn assign_to_windows<'a>( files: impl Iterator, time_window_size: i64, + mapper: &PrimaryKeyRangeMapper, ) -> BTreeMap { let mut windows: HashMap = HashMap::new(); // Iterates all files and assign to time windows according to max timestamp @@ -951,10 +1031,10 @@ fn assign_to_windows<'a>( match windows.entry(time_window) { Entry::Occupied(mut e) => { - e.get_mut().add_file(f.clone()); + e.get_mut().add_file(f.clone(), mapper); } Entry::Vacant(e) => { - let mut window = Window::new_with_file(f.clone()); + let mut window = Window::new_with_file(f.clone(), mapper); window.time_window = time_window; e.insert(window); } @@ -1064,7 +1144,6 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use bytes::Bytes; use common_base::Plugins; use common_time::range::TimestampRange; use store_api::storage::FileId; @@ -1075,7 +1154,8 @@ mod tests { use crate::compaction::test_util::{ compaction_region_with_ssts, new_file_handle, new_file_handle_with_sequence, new_file_handle_with_size_and_sequence, - new_file_handle_with_size_sequence_and_primary_key_range, + new_file_handle_with_size_sequence_and_primary_key_range, pk_range, + primary_key_mapper_for_test, }; use crate::config::MitoConfig; use crate::region::options::RegionOptions; @@ -1084,6 +1164,35 @@ mod tests { use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::scheduler_util::SchedulerEnv; + fn assign_to_windows<'a>( + files: impl Iterator, + time_window_size: i64, + ) -> BTreeMap { + super::assign_to_windows(files, time_window_size, &primary_key_mapper_for_test()) + } + + fn pick_count_first( + sorted_runs: Vec>, + max_output_file_size: Option, + ) -> Vec { + super::pick_count_first( + sorted_runs, + max_output_file_size, + &primary_key_mapper_for_test(), + ) + } + + fn pick_mixed_count_first( + sorted_runs: Vec>, + max_output_file_size: Option, + ) -> Vec { + super::pick_mixed_count_first( + sorted_runs, + max_output_file_size, + &primary_key_mapper_for_test(), + ) + } + impl TwcsPicker { async fn build_output_with_time_range( &self, @@ -1099,12 +1208,88 @@ mod tests { active_window, time_window_size, max_outputs: self.max_background_tasks, + primary_key_mapper: Arc::new(primary_key_mapper_for_test()), }, ) .await } } + #[test] + fn test_cross_schema_pk_overlap_before_window_aggregation() { + use crate::test_util::sst_util::{new_primary_key, sst_region_metadata}; + + let metadata = Arc::new(sst_region_metadata()); + let mut old = FileMeta { + region_id: metadata.region_id, + file_id: FileId::random(), + time_range: ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(1000), + ), + primary_key_min: Some(new_primary_key(&["a"]).into()), + primary_key_max: Some(new_primary_key(&["b"]).into()), + ..Default::default() + }; + let mut completed = new_primary_key(&["b"]); + completed.push(0); // The appended nullable tag's default. + let new = FileMeta { + file_id: FileId::random(), + time_range: ( + Timestamp::new_millisecond(500), + Timestamp::new_millisecond(2000), + ), + primary_key_min: Some(completed.clone().into()), + primary_key_max: Some(completed.clone().into()), + ..old.clone() + }; + assert!(old.primary_key_max < new.primary_key_min); + let mut ssts = SstVersion::new(metadata); + let old_id = old.file_id; + let new_id = new.file_id; + ssts.add_files( + crate::test_util::new_noop_file_purger(), + [old.clone(), new].into_iter(), + ); + let old_file = &ssts.levels()[0].files[&old_id]; + let new_file = &ssts.levels()[0].files[&new_id]; + let ranges = ssts.primary_key_mapper(); + let windows = super::assign_to_windows([old_file, new_file].into_iter(), 1, &ranges); + assert_eq!(2, windows.len()); + assert!( + windows + .values() + .all(|window| window_has_overlap(window, &windows)) + ); + + let mut window = Window::new_with_file(old_file.clone(), &ranges); + window.add_file(new_file.clone(), &ranges); + assert!(selected_overlaps_unselected( + std::slice::from_ref(new_file), + &window, + &ranges, + )); + assert_eq!( + Some(completed.into()), + window + .primary_key_range + .as_ref() + .map(|range| range.1.clone()) + ); + + // In the same target schema, a genuinely disjoint range still prunes. + old.file_id = FileId::random(); + old.primary_key_max = old.primary_key_min.clone(); + let disjoint_id = old.file_id; + ssts.add_files(crate::test_util::new_noop_file_purger(), [old].into_iter()); + let files = &ssts.levels()[0].files; + assert!(!files_overlap_inclusive( + &files[&disjoint_id], + &files[&new_id], + &ranges + )); + } + #[test] fn test_valid_max_input_files_env_overrides_default() { assert_eq!(64, parse_max_input_files(Some("64"))); @@ -1252,10 +1437,11 @@ mod tests { let env = SchedulerEnv::new().await; let metadata = metadata_for_test(); let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - let mut ssts = SstVersion::new(); + let mut ssts = SstVersion::new(metadata.clone()); ssts.add_files( Arc::new(crate::sst::file_purger::NoopFilePurger), [100, 101].into_iter().map(|sequence| FileMeta { + region_id: metadata.region_id, file_id: FileId::random(), time_range: ( Timestamp::new_millisecond(0), @@ -1530,10 +1716,6 @@ mod tests { /// (Window value, overlapping, files' time ranges in window) type ExpectedWindowSpec = (i64, bool, Vec<(i64, i64)>); - fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> { - Some((Bytes::from_static(min), Bytes::from_static(max))) - } - fn check_assign_to_windows_with_overlapping( file_time_ranges: &[(i64, i64)], time_window: i64, @@ -2554,6 +2736,7 @@ mod tests { l0_file_num: 8, l1_file_num: 2, phase: PickPhase::L1FileReduction, + primary_key_mapper: &primary_key_mapper_for_test(), }, None, ); @@ -2584,6 +2767,7 @@ mod tests { l0_file_num: 2, l1_file_num: 8, phase: PickPhase::L1FileReduction, + primary_key_mapper: &primary_key_mapper_for_test(), }, None, ); diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index b719543f9af..ab7ac0a5f9a 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -332,7 +332,7 @@ mod tests { let metadata = metadata_for_test(); let file_purger_ref = Arc::new(NoopFilePurger); - let mut ssts = SstVersion::new(); + let mut ssts = SstVersion::new(metadata.clone()); ssts.add_files( file_purger_ref, diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 547d3f213f9..092abc296a6 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -208,6 +208,7 @@ async fn test_strict_window_output_file_size( let version = engine.get_region(region_id).unwrap().version(); assert!(version.ssts.levels()[0].files.is_empty()); + let primary_key_mapper = version.ssts.primary_key_mapper(); let files = version.ssts.levels()[1].files().collect::>(); assert_eq!( 6000, @@ -231,7 +232,7 @@ async fn test_strict_window_output_file_size( } let mut ranges = window_files .iter() - .map(|file| file.primary_key_range().unwrap()) + .map(|file| file.primary_key_range(&primary_key_mapper).unwrap()) .collect::>(); ranges.sort_unstable(); assert!(ranges.windows(2).all(|pair| pair[0].1 < pair[1].0)); @@ -541,6 +542,313 @@ async fn assert_partial_compaction_preserves_delete_order(flat_format: bool) { ); } +/// Real-engine regression for the tombstone resurrection chain described in +/// https://greptime.feishu.cn/wiki/L2WYwIM40iG9tckowrnctyw8nQc. +/// No synthetic file sizes, PK ranges, or manually selected compaction inputs. +#[rstest::rstest] +#[tokio::test] +async fn test_cross_schema_compaction_keeps_rows_deleted( + #[values(false, true)] flat_format: bool, + #[values(false, true)] cross_window: bool, + #[values(None, Some(""))] tag_default: Option<&str>, +) { + use api::v1::SemanticType; + use api::v1::value::ValueData; + use datatypes::prelude::{ConcreteDataType, Value}; + use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema as DataColumnSchema}; + use store_api::metadata::ColumnMetadata; + use store_api::region_request::{AddColumn, AddColumnLocation, AlterKind}; + + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, mut columns) = env_for_manual_compaction_with_window( + &mut env, + region_id, + flat_format, + if cross_window { "2h" } else { "1h" }, + ) + .await; + let (target_ts, other_ts) = if cross_window { (3500, 3700) } else { (10, 1) }; + + // Incompressible, distinct a* keys keep the old SST large enough for a + // partial pick to leave it behind. Its maximum PK is exactly the victim b. + let mut old_rows = build_rows_for_key("b", target_ts, target_ts + 1, 0); + let mut rng = ::seed_from_u64(0); + for _ in 0..3000 { + old_rows.extend(build_rows_for_key( + &format!("a{:032x}", rand::Rng::random::(&mut rng)), + other_ts, + other_ts + 1, + 0, + )); + } + put_rows( + &engine, + region_id, + Rows { + schema: columns.clone(), + rows: old_rows, + }, + ) + .await; + flush(&engine, region_id).await; + let old_version = engine.get_region(region_id).unwrap().version(); + let old_files = old_version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .collect::>(); + assert_eq!(1, old_files.len()); + let old_file = old_files[0]; + if cross_window { + // This real SST crosses the new one-hour boundary. TWCS assigns it by + // max_ts=3700 to window 7200, unlike Delete(b,3500) in window 3600. + set_compaction_window(&engine, region_id, "1h").await; + } + + let added = ColumnMetadata { + column_id: 3, + semantic_type: SemanticType::Tag, + column_schema: DataColumnSchema::new("tag_1", ConcreteDataType::string_datatype(), true) + .with_default_constraint( + tag_default.map(|value| ColumnDefaultConstraint::Value(Value::from(value))), + ) + .unwrap(), + }; + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: added.clone(), + location: Some(AddColumnLocation::First), + }], + }, + }), + ) + .await + .unwrap(); + columns.push(column_metadata_to_column_schema(&added)); + let append_tag = |mut rows: Vec| { + for row in &mut rows { + row.values.push(api::v1::Value { + value_data: tag_default.map(|value| ValueData::StringValue(value.to_string())), + }); + } + Rows { + schema: columns.clone(), + rows, + } + }; + let deletes = ["b", "c"] + .into_iter() + .flat_map(|key| build_rows_for_key(key, target_ts, target_ts + 1, 0)) + .collect(); + engine + .handle_request( + region_id, + RegionRequest::Delete(RegionDeleteRequest { + rows: append_tag(deletes), + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + flush(&engine, region_id).await; + for ts in target_ts + 10..target_ts + 25 { + put_rows( + &engine, + region_id, + append_tag(build_rows_for_key("c", ts, ts + 1, 0)), + ) + .await; + flush(&engine, region_id).await; + } + + let region = engine.get_region(region_id).unwrap(); + let table_dir = region.table_dir().to_string(); + let version = region.version(); + let files = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .collect::>(); + assert_eq!(17, files.len()); + assert_eq!(vec![0, 3], version.metadata.primary_key); + let new_ids = files + .iter() + .filter(|file| file.file_id() != old_file.file_id()) + .map(|file| file.file_id()) + .collect::>(); + for file in files + .iter() + .filter(|file| new_ids.contains(&file.file_id())) + { + assert!(old_file.meta_ref().primary_key_max < file.meta_ref().primary_key_min); + assert!(file.time_range().1 < Timestamp::new_second(3600)); + } + assert_eq!( + cross_window, + old_file.time_range().1 > Timestamp::new_second(3600) + ); + + let picked = preview_regular_compaction(&engine, region_id).await; + assert_eq!(1, picked.outputs.len()); + let output = &picked.outputs[0]; + assert_eq!( + new_ids, + output.inputs.iter().map(|file| file.file_id()).collect() + ); + let filter_deleted = output.filter_deleted; + let before = sorted_scan_timestamps(&engine, region_id).await; + assert_eq!(3015, before.len()); + assert!(!before.contains(&(target_ts as i64 * 1000))); + + compact(&engine, region_id).await; + + let after_version = region.version(); + let after_files = after_version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .collect::>(); + assert_eq!(2, after_files.len()); + assert!( + after_files + .iter() + .any(|file| file.file_id() == old_file.file_id()) + ); + let output_file = after_files + .into_iter() + .find(|file| file.file_id() != old_file.file_id()) + .unwrap(); + assert!(!new_ids.contains(&output_file.file_id())); + let output_rows = output_file.meta_ref().num_rows; + let output_deletes = count_sst_deletes(&engine, region_id, output_file).await; + let after = sorted_scan_timestamps(&engine, region_id).await; + + let config = MitoConfig { + default_flat_format: flat_format, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.reopen_engine(engine, config).await; + crate::test_util::reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let reopened = sorted_scan_timestamps(&engine, region_id).await; + // Do not fail on filter_deleted before actually executing compaction/reopen: + // the negative-control run must demonstrate data resurrection, not just a bad plan. + assert!( + !filter_deleted + && output_rows == 17 + && output_deletes == 2 + && after == before + && reopened == before, + "cross-schema tombstone loss: flat_format={flat_format}, cross_window={cross_window}, default={tag_default:?}, selected={}, filter_deleted={filter_deleted}, output_rows={output_rows}, output_deletes={output_deletes}, before={}, after={}, reopened={}, resurrected_after={}, resurrected_reopen={}", + new_ids.len(), + before.len(), + after.len(), + reopened.len(), + after.contains(&(target_ts as i64 * 1000)), + reopened.contains(&(target_ts as i64 * 1000)), + ); +} + +async fn count_sst_deletes( + engine: &MitoEngine, + region_id: RegionId, + file: &crate::sst::file::FileHandle, +) -> usize { + use api::v1::OpType; + use datatypes::arrow::datatypes::UInt8Type; + + let region = engine.get_region(region_id).unwrap(); + let mut reader = region + .access_layer + .read_sst(file.clone()) + .build() + .await + .unwrap() + .unwrap(); + let mut deletes = 0; + while let Some(batch) = reader.next_record_batch().await.unwrap() { + deletes += batch + .column(batch.num_columns() - 1) + .as_primitive::() + .values() + .iter() + .filter(|op| **op == OpType::Delete as u8) + .count(); + } + deletes +} + +async fn set_compaction_window(engine: &MitoEngine, region_id: RegionId, window: &str) { + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: SetRegionOptions { + options: vec![SetRegionOption::Twsc( + "compaction.twcs.time_window".into(), + window.into(), + )], + }, + }), + ) + .await + .unwrap(); +} + +async fn sorted_scan_timestamps(engine: &MitoEngine, region_id: RegionId) -> Vec { + let stream = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .scan() + .await + .unwrap(); + let mut timestamps = collect_stream_ts(stream).await; + timestamps.sort_unstable(); + timestamps +} + +async fn preview_regular_compaction( + engine: &MitoEngine, + region_id: RegionId, +) -> crate::compaction::picker::PickerOutput { + use crate::compaction::compactor::CompactionRegion; + use crate::compaction::picker::new_picker; + + let region = engine.get_region(region_id).unwrap(); + let version = region.version(); + let picker = new_picker( + &RegionCompactRequest::default().options, + &version.options, + None, + None, + ); + let compaction_region = CompactionRegion { + region_id, + region_options: version.options.clone(), + engine_config: Arc::new(MitoConfig::default()), + region_metadata: version.metadata.clone(), + cache_manager: engine.cache_manager(), + access_layer: region.access_layer.clone(), + manifest_ctx: region.manifest_ctx.clone(), + current_version: version.into(), + file_purger: None, + ttl: None, + max_parallelism: 1, + plugins: common_base::Plugins::new(), + }; + picker.pick(&compaction_region).await.unwrap().unwrap() +} + struct CompactionListenerGuard(Option>); impl CompactionListenerGuard { @@ -1790,6 +2098,15 @@ async fn env_for_manual_compaction( env: &mut TestEnv, region_id: RegionId, flat_format: bool, +) -> (MitoEngine, Vec) { + env_for_manual_compaction_with_window(env, region_id, flat_format, "1h").await +} + +async fn env_for_manual_compaction_with_window( + env: &mut TestEnv, + region_id: RegionId, + flat_format: bool, + time_window: &str, ) -> (MitoEngine, Vec) { let engine = env .create_engine(MitoConfig { @@ -1812,7 +2129,7 @@ async fn env_for_manual_compaction( let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") - .insert_option("compaction.twcs.time_window", "1h") + .insert_option("compaction.twcs.time_window", time_window) .build(); let column_schemas = request .column_metadatas diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 21416bf3652..1592bfebd60 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1138,6 +1138,21 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid SST primary key range: {reason}"))] + InvalidPrimaryKeyRange { + reason: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to decode SST primary key range {endpoint} endpoint"))] + DecodePrimaryKeyRange { + endpoint: &'static str, + source: mito_codec::error::Error, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Region {} is busy", region_id))] RegionBusy { region_id: RegionId, @@ -1586,7 +1601,10 @@ impl ErrorExt for Error { FulltextPushText { source, .. } | FulltextFinish { source, .. } | ApplyFulltextIndex { source, .. } => source.status_code(), - DecodeStats { .. } | StatsNotPresent { .. } => StatusCode::Internal, + DecodeStats { .. } + | StatsNotPresent { .. } + | InvalidPrimaryKeyRange { .. } + | DecodePrimaryKeyRange { .. } => StatusCode::Internal, RegionBusy { .. } => StatusCode::RegionBusy, GetSchemaMetadata { source, .. } => source.status_code(), Timeout { .. } => StatusCode::Cancelled, diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 7c93d137c04..1d0991a6bfb 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -17,7 +17,7 @@ use std::collections::{BTreeMap, HashSet}; use std::fmt; use std::num::NonZeroU64; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::time::Instant; use api::v1::SemanticType; @@ -90,6 +90,7 @@ use crate::sst::index::vector_index::applier::{VectorIndexApplier, VectorIndexAp use crate::sst::parquet::Json2RewriteTargets; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::reader::ReaderMetrics; +use crate::sst::primary_key::PrimaryKeyRangeMapper; #[cfg(feature = "vector_index")] const VECTOR_INDEX_OVERFETCH_MULTIPLIER: usize = 2; @@ -589,6 +590,7 @@ impl ScanRegion { .with_predicate(predicate) .with_memtables(mem_range_builders) .with_files(files) + .with_primary_key_mapper(self.version.ssts.primary_key_mapper()) .with_cache(self.cache_strategy) .with_inverted_index_appliers(inverted_index_appliers) .with_bloom_filter_index_appliers(bloom_filter_appliers) @@ -977,6 +979,8 @@ pub struct ScanInput { pub(crate) memtables: Vec, /// Handles to SST files to scan. pub(crate) files: Vec, + /// Shares the pinned schema's encoded defaults across parallel range readers. + primary_key_mapper: OnceLock>, /// Scan-wide hint for rows in an execution batch. batch_size: usize, /// Cache. @@ -1066,6 +1070,7 @@ impl ScanInput { region_partition_expr: None, memtables: Vec::new(), files: Vec::new(), + primary_key_mapper: OnceLock::new(), batch_size: crate::sst::parquet::DEFAULT_READ_BATCH_SIZE, cache_strategy: CacheStrategy::Disabled, ignore_file_not_found: false, @@ -1101,6 +1106,12 @@ impl ScanInput { self.batch_size } + /// Interprets file statistics using this scan's pinned schema. + pub(crate) fn primary_key_mapper(&self) -> &PrimaryKeyRangeMapper { + self.primary_key_mapper + .get_or_init(|| Arc::new(PrimaryKeyRangeMapper::new(self.region_metadata().clone()))) + } + /// Returns the range implied by the range-cache time filters. pub(crate) fn implied_time_range(&self) -> Option<&TimestampRange> { self.scan_analysis @@ -1121,6 +1132,11 @@ impl ScanInput { } impl ScanInputBuilder { + fn with_primary_key_mapper(mut self, mapper: Arc) -> Self { + self.input.primary_key_mapper = OnceLock::from(mapper); + self + } + /// Sets whether to ignore range indexes during scans. #[must_use] pub(crate) fn with_ignore_range_index(mut self, ignore: bool) -> Self { diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 9e86ff1e5e9..5ac69d191a6 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -408,7 +408,7 @@ async fn build_series_partition_range( if stream_ctx.is_file_range_index(index) { let file = stream_ctx.input.file_from_index(index); if matches!( - file.primary_key_range() + file.primary_key_range(stream_ctx.input.primary_key_mapper()) .and_then(|(min, max)| { filter.overlaps_encoded_bounds(&codec, &min, &max) }), Some(false) ) { diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 033f9916905..fb5d260fe6f 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1232,7 +1232,7 @@ async fn preload_parquet_meta_cache_for_files( .get_compact_sst_meta_data(file_id, PageIndexPolicy::Optional) .await { - if file_handle.primary_key_range().is_none() + if file_handle.raw_primary_key_range().is_none() && let Some(primary_key_range) = extract_primary_key_range( metadata.parquet_metadata().as_ref(), ®ion_metadata, @@ -1251,7 +1251,7 @@ async fn preload_parquet_meta_cache_for_files( .await { let decoded = metadata.decoded(); - if file_handle.primary_key_range().is_none() + if file_handle.raw_primary_key_range().is_none() && let Some(primary_key_range) = extract_primary_key_range( decoded.parquet_metadata().as_ref(), ®ion_metadata, @@ -1664,7 +1664,9 @@ mod tests { let file_id = FileId::random(); let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; - let primary_key = Arc::new(BinaryArray::from_iter_values([b"a", b"b", b"c"])) as ArrayRef; + let keys = + ["a", "b", "c"].map(|tag| crate::test_util::sst_util::new_primary_key(&[tag, ""])); + let primary_key = Arc::new(BinaryArray::from_iter_values(&keys)) as ArrayRef; let batch = RecordBatch::try_from_iter([ ("col", col), ( @@ -1744,7 +1746,7 @@ mod tests { .await .is_some() ); - assert!(file_handle.primary_key_range().is_some()); + assert!(file_handle.raw_primary_key_range().is_some()); } #[tokio::test] diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index e9ba4a1cdbf..a56cef1f544 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -404,9 +404,9 @@ impl VersionBuilder { /// Returns a new builder. pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitionsRef) -> Self { VersionBuilder { - metadata, + metadata: metadata.clone(), memtables: Arc::new(MemtableVersion::new(mutable)), - ssts: Arc::new(SstVersion::new()), + ssts: Arc::new(SstVersion::new(metadata)), flushed_entry_id: 0, flushed_sequence: 0, truncated_entry_id: None, @@ -437,6 +437,9 @@ impl VersionBuilder { /// Sets metadata. pub(crate) fn metadata(mut self, metadata: RegionMetadataRef) -> Self { + if !Arc::ptr_eq(&self.metadata, &metadata) { + Arc::make_mut(&mut self.ssts).set_metadata(metadata.clone()); + } self.metadata = metadata; self } @@ -525,7 +528,7 @@ impl VersionBuilder { /// Clear all files in the builder. pub(crate) fn clear_files(mut self) -> Self { - self.ssts = Arc::new(SstVersion::new()); + self.ssts = Arc::new(SstVersion::new(self.metadata.clone())); self } diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index e53c28e54c7..7a0e3603836 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -47,6 +47,7 @@ pub mod file_ref; pub mod index; pub mod location; pub mod parquet; +pub(crate) mod primary_key; pub mod range_index; pub(crate) mod version; diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 9702477b07a..f160bbb3785 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -24,7 +24,7 @@ use std::sync::{Arc, Mutex, RwLock}; use base64::prelude::{BASE64_STANDARD, Engine}; use bytes::Bytes; use common_base::readable_size::ReadableSize; -use common_telemetry::{debug, error}; +use common_telemetry::{debug, error, warn}; use common_time::Timestamp; use partition::expr::PartitionExpr; use serde::{Deserialize, Serialize}; @@ -39,6 +39,7 @@ use crate::cache::file_cache::{FileType, IndexKey}; use crate::sst::file_purger::FilePurgerRef; use crate::sst::location; use crate::sst::parquet::SstInfo; +use crate::sst::primary_key::PrimaryKeyRangeMapper; /// Custom serde functions for Bytes fields serialized as base64 strings. fn serialize_bytes_option(bytes: &Option, serializer: S) -> Result @@ -497,6 +498,7 @@ impl fmt::Debug for FileHandle { } impl FileHandle { + /// Creates a handle sharing the file's physical state and original statistics. pub fn new(meta: FileMeta, file_purger: FilePurgerRef) -> FileHandle { let pk_range = meta.primary_key_range(); FileHandle { @@ -612,12 +614,100 @@ impl FileHandle { self.inner.deleted.load(Ordering::Relaxed) } - pub fn primary_key_range(&self) -> Option<(Bytes, Bytes)> { - self.inner.primary_key_range.read().unwrap().clone() + /// Returns bounds aligned to the caller's pinned schema, before any comparison + /// or aggregation. Unknown and invalid statistics cannot exclude possible data. + pub(crate) fn primary_key_range( + &self, + mapper: &PrimaryKeyRangeMapper, + ) -> Option<(Bytes, Bytes)> { + debug_assert_eq!(self.region_id().table_id(), mapper.region_id().table_id()); + if let Some(range) = self + .inner + .primary_key_range + .read() + .unwrap() + .aligned(mapper.schema_version()) + { + return range.clone(); + } + // Recheck under the write lock: another snapshot may have replaced the cached schema. + let aligned = self.inner.primary_key_range.write().unwrap().align(mapper); + match aligned { + Ok(range) => range, + Err(err) => { + warn!(err; "Invalid SST primary key range; using unknown bounds, region: {}, file: {}, schema version: {}", + self.region_id(), self.file_id(), mapper.schema_version()); + None + } + } + } + + /// Returns original statistics for metadata hydration, never schema-aligned bounds. + pub fn raw_primary_key_range(&self) -> Option<(Bytes, Bytes)> { + self.inner.primary_key_range.read().unwrap().raw().cloned() } pub(crate) fn set_primary_key_range(&self, primary_key_range: (Bytes, Bytes)) { - *self.inner.primary_key_range.write().unwrap() = Some(primary_key_range); + // SST contents are immutable. Hydrate missing raw statistics without + // replacing the source of already cached schema views. + let mut range = self.inner.primary_key_range.write().unwrap(); + if matches!(*range, PrimaryKeyRange::Missing) { + *range = PrimaryKeyRange::Raw(primary_key_range); + } + } +} + +type PrimaryKeyBounds = (Bytes, Bytes); + +/// A single-slot cache shared by file handles. Always retain the source bounds so +/// older snapshots and changed defaults can realign without interpreting padded values as stored. +enum PrimaryKeyRange { + Missing, + Raw(PrimaryKeyBounds), + Aligned { + raw: PrimaryKeyBounds, + schema_version: u64, + bounds: Option, + }, +} + +impl PrimaryKeyRange { + fn raw(&self) -> Option<&PrimaryKeyBounds> { + match self { + Self::Missing => None, + Self::Raw(raw) | Self::Aligned { raw, .. } => Some(raw), + } + } + + fn aligned(&self, target_version: u64) -> Option<&Option> { + match self { + Self::Aligned { + schema_version, + bounds, + .. + } if *schema_version == target_version => Some(bounds), + _ => None, + } + } + + fn align( + &mut self, + mapper: &PrimaryKeyRangeMapper, + ) -> crate::error::Result> { + if let Some(bounds) = self.aligned(mapper.schema_version()) { + return Ok(bounds.clone()); + } + let Some(raw) = self.raw().cloned() else { + return Ok(None); + }; + let aligned = mapper.map(raw.clone()); + // Failed mappings also occupy the cache, so repeated hits don't repeat the warning. + *self = Self::Aligned { + raw, + schema_version: mapper.schema_version(), + bounds: aligned.as_ref().ok().cloned().flatten(), + }; + aligned } } @@ -629,7 +719,7 @@ struct FileHandleInner { compacting: AtomicBool, deleted: AtomicBool, index_outdated: AtomicBool, - primary_key_range: RwLock>, + primary_key_range: RwLock, file_purger: FilePurgerRef, } @@ -656,7 +746,9 @@ impl FileHandleInner { compacting: AtomicBool::new(false), deleted: AtomicBool::new(false), index_outdated: AtomicBool::new(false), - primary_key_range: RwLock::new(primary_key_range), + primary_key_range: RwLock::new( + primary_key_range.map_or(PrimaryKeyRange::Missing, PrimaryKeyRange::Raw), + ), file_purger, } } diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index 5ff88c0874b..1b3b6239897 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -71,15 +71,15 @@ impl fmt::Debug for LocalFilePurger { } } -#[cfg(not(debug_assertions))] +#[cfg(not(any(debug_assertions, test)))] /// Whether to enable GC for the file purger. pub fn should_enable_gc(global_gc_enabled: bool, object_store_scheme: &'static str) -> bool { global_gc_enabled && object_store_scheme != object_store::services::FS_SCHEME } -#[cfg(debug_assertions)] -/// For debug build, we may use Fs as the object store scheme, -/// so we need to enable GC for local file system. +#[cfg(any(debug_assertions, test))] +/// Debug builds and unit tests use Fs to exercise object-store GC, including +/// unit tests compiled with the release profile. pub fn should_enable_gc(global_gc_enabled: bool, _object_store_scheme: &'static str) -> bool { global_gc_enabled } diff --git a/src/mito2/src/sst/parquet/metadata.rs b/src/mito2/src/sst/parquet/metadata.rs index 3c0d197c5c6..8e6a0a5998c 100644 --- a/src/mito2/src/sst/parquet/metadata.rs +++ b/src/mito2/src/sst/parquet/metadata.rs @@ -144,6 +144,11 @@ pub(crate) fn extract_primary_key_range( return None; }; + // Truncated bounds can end at a field boundary and masquerade as an + // older Dense schema. Only complete endpoints can be schema-normalized. + if !stats.min_is_exact() || !stats.max_is_exact() { + return None; + } let row_group_min = Bytes::copy_from_slice(stats.min_bytes_opt()?); let row_group_max = Bytes::copy_from_slice(stats.max_bytes_opt()?); min = Some(match min { @@ -321,6 +326,16 @@ mod tests { ); } + #[test] + fn test_extract_primary_key_range_rejects_truncated_statistics() { + let key = vec![b'a'; 1024]; + let metadata = build_test_metadata(true, &[&key], &[1], EnabledStatistics::Page); + assert_eq!( + None, + extract_primary_key_range(&metadata, &sst_region_metadata()) + ); + } + #[test] fn test_extract_primary_key_range_returns_none_when_any_rg_stats_missing() { let metadata = build_test_metadata( diff --git a/src/mito2/src/sst/primary_key.rs b/src/mito2/src/sst/primary_key.rs new file mode 100644 index 00000000000..5fb258052c7 --- /dev/null +++ b/src/mito2/src/sst/primary_key.rs @@ -0,0 +1,847 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Schema-bound views of persisted primary key ranges. + +use std::sync::OnceLock; + +use bytes::Bytes; +use datatypes::prelude::ConcreteDataType; +use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodec, SortField}; +use snafu::{ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; +use store_api::storage::RegionId; + +use crate::error::{DecodePrimaryKeyRangeSnafu, InvalidPrimaryKeyRangeSnafu, Result}; + +/// Schema conversion and default encodings shared by a version's comparison contexts. +#[derive(Debug)] +pub(crate) struct PrimaryKeyRangeMapper { + /// Target metadata pinned by the owning version or scan, not the SST's write schema. + /// ALTER creates a new mapper; existing snapshots keep their original target schema. + metadata: RegionMetadataRef, + codec: DensePrimaryKeyCodec, + /// Encoding a constant is independent of the file's missing-prefix length. + encoded_defaults: Vec>>, + suffixes: Vec>>, +} + +impl PrimaryKeyRangeMapper { + pub(crate) fn new(metadata: RegionMetadataRef) -> Self { + let codec = DensePrimaryKeyCodec::new(&metadata); + let encoded_defaults = (0..codec.num_fields()).map(|_| OnceLock::new()).collect(); + let suffixes = (0..codec.num_fields()).map(|_| OnceLock::new()).collect(); + Self { + metadata, + codec, + encoded_defaults, + suffixes, + } + } + + /// Reuses encoded constants without retaining a chain of old schema snapshots. + pub(crate) fn with_metadata(&self, metadata: RegionMetadataRef) -> Self { + let mut mapper = Self::new(metadata); + for (index, (old, new)) in self + .metadata + .primary_key_columns() + .zip(mapper.metadata.primary_key_columns()) + .enumerate() + { + if same_pk_column(old, new) { + mapper.encoded_defaults[index] = self.encoded_defaults[index].clone(); + } + } + mapper + } + + /// Returns the schema version used to interpret the bounds. + pub(crate) fn schema_version(&self) -> u64 { + self.metadata.schema_version + } + + /// Returns the region id of mapper. + pub(crate) fn region_id(&self) -> RegionId { + self.metadata.region_id + } + + /// Maps bounds from the same table's encoding and a compatible schema prefix. + /// Returns unknown for defaults that cannot be safely encoded. + /// Invalid bounds return an error so callers can diagnose them before falling back + /// to unknown. Neither the error nor its source includes the encoded keys. + pub(crate) fn map(&self, (min, max): (Bytes, Bytes)) -> Result> { + ensure!( + min <= max, + InvalidPrimaryKeyRangeSnafu { + reason: "min is greater than max", + } + ); + // Sparse-to-sparse read compatibility preserves encoded keys verbatim. + if self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse { + return Ok(Some((min, max))); + } + // Ordinary same-table Dense ALTER only appends PK fields and preserves + // prefix order/types. SyncColumns violating that invariant is out of scope. + let prefix_len = self.range_prefix_len(&min, &max)?; + if prefix_len == self.codec.num_fields() { + return Ok(Some((min, max))); + } + let Some(suffix) = self.suffixes[prefix_len] + .get_or_init(|| self.encode_suffix(prefix_len)) + .as_ref() + else { + return Ok(None); + }; + // Fixed-layout Dense keys are prefix-free, so appending one constant + // suffix is monotone. Transform each file before aggregating any ranges. + Ok(Some(( + append_suffix(min, suffix), + append_suffix(max, suffix), + ))) + } + + fn range_prefix_len(&self, min: &[u8], max: &[u8]) -> Result { + let min_len = self + .codec + .decode_prefix_len(min) + .context(DecodePrimaryKeyRangeSnafu { endpoint: "min" })?; + let max_len = self + .codec + .decode_prefix_len(max) + .context(DecodePrimaryKeyRangeSnafu { endpoint: "max" })?; + ensure!( + min_len == max_len, + InvalidPrimaryKeyRangeSnafu { + reason: format!( + "endpoints have different field counts: min {min_len}, max {max_len}" + ), + } + ); + Ok(min_len) + } + + fn encode_suffix(&self, prefix_len: usize) -> Option { + let mut suffix = Vec::with_capacity(self.codec.num_fields() - prefix_len); + for (index, column) in self + .metadata + .primary_key_columns() + .enumerate() + .skip(prefix_len) + { + let encoded = self.encoded_defaults[index] + .get_or_init(|| encode_default(column)) + .as_ref()?; + suffix.extend_from_slice(encoded); + } + Some(suffix.into()) + } +} + +fn same_pk_column(left: &ColumnMetadata, right: &ColumnMetadata) -> bool { + left.column_id == right.column_id + && left.column_schema.data_type == right.column_schema.data_type + && left.column_schema.is_nullable() == right.column_schema.is_nullable() + && left.column_schema.default_constraint() == right.column_schema.default_constraint() +} + +/// Missing, impure or unencodable defaults cannot produce exact bounds. +fn encode_default(column: &ColumnMetadata) -> Option { + let schema = &column.column_schema; + if schema.is_default_impure() { + return None; + } + let default = schema.create_default().ok()??; + let field = SortField::new(schema.data_type.clone()); + if default.is_null() { + // Dense uses one NULL marker per field, but unsupported types must stay unknown. + return match field.encode_data_type() { + ConcreteDataType::Null(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) => None, + _ => Some(Bytes::from_static(&[0])), + }; + } + let mut encoded = Vec::new(); + DensePrimaryKeyCodec::with_fields(vec![(column.column_id, field)]) + .encode_values(&[(column.column_id, default)], &mut encoded) + .ok()?; + Some(encoded.into()) +} + +fn append_suffix(key: Bytes, suffix: &Bytes) -> Bytes { + let mut completed = Vec::with_capacity(key.len() + suffix.len()); + completed.extend_from_slice(&key); + completed.extend_from_slice(suffix); + completed.into() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::v1::SemanticType; + use common_time::Timestamp; + use datatypes::prelude::{ConcreteDataType, Value}; + use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; + use rstest::rstest; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::storage::FileId; + + use super::*; + use crate::compaction::run::files_overlap_inclusive; + use crate::manifest::action::RegionEdit; + use crate::memtable::time_partition::TimePartitions; + use crate::memtable::time_series::TimeSeriesMemtableBuilder; + use crate::region::version::{VersionBuilder, VersionRef}; + use crate::sst::file::{FileHandle, FileMeta}; + use crate::test_util::new_noop_file_purger; + + fn metadata(defaults: &[Value]) -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + for (id, value) in defaults.iter().enumerate() { + let data_type = if value.is_null() { + ConcreteDataType::string_datatype() + } else { + value.data_type() + }; + builder.push_column_metadata(ColumnMetadata { + column_id: id as u32, + semantic_type: SemanticType::Tag, + column_schema: ColumnSchema::new(format!("tag_{id}"), data_type, true) + .with_default_constraint(Some(ColumnDefaultConstraint::Value(value.clone()))) + .unwrap(), + }); + } + builder.push_column_metadata(ColumnMetadata { + column_id: 100, + semantic_type: SemanticType::Timestamp, + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + }); + builder.primary_key((0..defaults.len() as u32).collect()); + Arc::new(builder.build().unwrap()) + } + + fn metadata_at_version(defaults: &[Value], version: u64) -> RegionMetadataRef { + let mut metadata = (*metadata(defaults)).clone(); + metadata.schema_version = version; + Arc::new(metadata) + } + + fn encode(metadata: &RegionMetadataRef, values: &[Value]) -> Bytes { + let mut bytes = Vec::new(); + let values: Vec<_> = values + .iter() + .enumerate() + .map(|(id, value)| (id as u32, value.clone())) + .collect(); + DensePrimaryKeyCodec::new(metadata) + .encode_values(&values, &mut bytes) + .unwrap(); + bytes.into() + } + + fn file_meta(metadata: &RegionMetadataRef, values: &[Value]) -> FileMeta { + let key = encode(metadata, values); + FileMeta { + region_id: metadata.region_id, + file_id: FileId::random(), + time_range: (Timestamp::new_millisecond(0), Timestamp::new_millisecond(1)), + primary_key_min: Some(key.clone()), + primary_key_max: Some(key), + ..Default::default() + } + } + + fn version_with_files(metadata: RegionMetadataRef, files: Vec) -> VersionRef { + let mutable = Arc::new(TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + None, + )); + Arc::new( + VersionBuilder::new(metadata, mutable) + .add_files(new_noop_file_purger(), files.into_iter()) + .build(), + ) + } + + #[test] + fn test_field_addition_updates_cache_version_and_shares_sst_list() { + let metadata = metadata(&[Value::from(""), Value::from("default")]); + let raw = file_meta(&metadata, &[Value::from("a")]); + let file_id = raw.file_id; + let original = version_with_files(metadata.clone(), vec![raw]); + let mapper = original.ssts.primary_key_mapper(); + let range = original.ssts.levels()[0].files[&file_id] + .primary_key_range(&mapper) + .unwrap(); + let mut builder = RegionMetadataBuilder::from_existing((*metadata).clone()); + builder + .push_column_metadata(ColumnMetadata { + column_id: 101, + semantic_type: SemanticType::Field, + column_schema: ColumnSchema::new("field", ConcreteDataType::int64_datatype(), true), + }) + .bump_version(); + let changed = VersionBuilder::from_version(original.clone()) + .metadata(Arc::new(builder.build().unwrap())) + .build(); + + assert_eq!(1, changed.metadata.schema_version); + assert!(changed.metadata.column_by_id(101).is_some()); + assert_eq!( + original.ssts.levels().as_ptr(), + changed.ssts.levels().as_ptr() + ); + let changed_mapper = changed.ssts.primary_key_mapper(); + assert_eq!(1, changed_mapper.schema_version()); + let changed_range = changed.ssts.levels()[0].files[&file_id] + .primary_key_range(&changed_mapper) + .unwrap(); + assert_eq!(range, changed_range); + let cached = changed.ssts.levels()[0].files[&file_id] + .primary_key_range(&changed_mapper) + .unwrap(); + assert_eq!(cached.0.as_ptr(), changed_range.0.as_ptr()); + assert_eq!(cached.1.as_ptr(), changed_range.1.as_ptr()); + } + + // SET/DROP DEFAULT may change padding even when the number of PK columns stays fixed. + #[rstest] + #[case::set(Value::from("x"), Some(Value::from("y")))] + #[case::set_null(Value::from("x"), Some(Value::Null))] + #[case::drop(Value::from("x"), None)] + #[case::replace_null(Value::Null, Some(Value::from("y")))] + fn test_tag_default_changes_refresh_only_missing_values( + #[case] previous_default: Value, + #[case] new_default: Option, + ) { + let metadata = metadata(&[Value::from(""), previous_default.clone()]); + let missing = file_meta(&metadata, &[Value::from("a")]); + let complete = file_meta(&metadata, &[Value::from("b"), Value::from("stored")]); + let missing_id = missing.file_id; + let complete_id = complete.file_id; + let complete_range = complete.primary_key_range(); + let original = version_with_files(metadata.clone(), vec![missing, complete]); + let old_mapper = original.ssts.primary_key_mapper(); + let old_range = original.ssts.levels()[0].files[&missing_id].primary_key_range(&old_mapper); + + let mut changed_metadata = (*metadata).clone(); + changed_metadata.column_metadatas[1].column_schema = changed_metadata.column_metadatas[1] + .column_schema + .clone() + .with_default_constraint(new_default.clone().map(ColumnDefaultConstraint::Value)) + .unwrap(); + let mut builder = RegionMetadataBuilder::from_existing(changed_metadata); + builder.bump_version(); + let changed_metadata = Arc::new(builder.build().unwrap()); + let changed = VersionBuilder::from_version(original.clone()) + .metadata(changed_metadata.clone()) + .build(); + + let expected = encode( + &changed_metadata, + &[Value::from("a"), new_default.unwrap_or(Value::Null)], + ); + let changed_mapper = changed.ssts.primary_key_mapper(); + assert_eq!( + Some((expected.clone(), expected)), + changed.ssts.levels()[0].files[&missing_id].primary_key_range(&changed_mapper) + ); + assert_eq!( + complete_range, + changed.ssts.levels()[0].files[&complete_id].primary_key_range(&changed_mapper) + ); + let expected_old = encode(&metadata, &[Value::from("a"), previous_default]); + assert_eq!(Some((expected_old.clone(), expected_old)), old_range); + assert_eq!( + old_range, + original.ssts.levels()[0].files[&missing_id].primary_key_range(&old_mapper) + ); + } + + #[rstest] + #[case::mixed(vec![ + Value::from("default"), + Value::Null, + Value::Int64(42), + Value::Binary(vec![0, 255].into()), + ])] + #[case::nulls(vec![Value::Null; 4])] + #[case::empty(vec![])] + fn test_dense_ranges_complete_every_historical_prefix(#[case] defaults: Vec) { + let metadata = metadata(&defaults); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + let expected = encode(&metadata, &defaults); + for count in 0..=defaults.len() { + let prefix = encode(&metadata, &defaults[..count]); + assert_eq!( + Some((expected.clone(), expected.clone())), + mapper.map((prefix.clone(), prefix)).unwrap() + ); + } + } + + #[rstest] + #[case::cold(false)] + #[case::warm(true)] + fn test_successive_pk_appends_reuse_constants_and_complete_all_prefixes(#[case] warm: bool) { + let defaults = [ + Value::from(""), + Value::from("constant"), + Value::Null, + Value::Int64(42), + Value::Null, + ]; + let mut mapper = PrimaryKeyRangeMapper::new(metadata(&defaults[..2])); + let mut values = defaults.clone(); + values[0] = Value::from("stored"); + let raw = encode(&mapper.metadata, &values[..1]); + if warm { + assert!(mapper.map((raw.clone(), raw)).unwrap().is_some()); + } + + for count in 3..=defaults.len() { + let next_metadata = metadata(&defaults[..count]); + let next_mapper = mapper.with_metadata(next_metadata.clone()); + if let Some(Some(encoded)) = mapper.encoded_defaults[1].get() { + // An already encoded non-NULL constant must share its allocation across ALTER. + let reused = next_mapper.encoded_defaults[1] + .get() + .unwrap() + .as_ref() + .unwrap(); + assert_eq!(encoded.as_ptr(), reused.as_ptr()); + } + let expected = encode(&next_metadata, &values[..count]); + for prefix_len in 1..=count { + let prefix = encode(&next_metadata, &values[..prefix_len]); + assert_eq!( + Some((expected.clone(), expected.clone())), + next_mapper.map((prefix.clone(), prefix)).unwrap() + ); + } + let old_expected = encode(&mapper.metadata, &values[..count - 1]); + let raw = encode(&mapper.metadata, &values[..1]); + assert_eq!( + Some((old_expected.clone(), old_expected)), + mapper.map((raw.clone(), raw)).unwrap() + ); + mapper = next_mapper; + } + } + + #[rstest] + #[case::string(ConcreteDataType::string_datatype(), true, true)] + #[case::int(ConcreteDataType::int64_datatype(), true, true)] + #[case::binary(ConcreteDataType::binary_datatype(), true, true)] + #[case::timestamp(ConcreteDataType::timestamp_millisecond_datatype(), true, true)] + #[case::required(ConcreteDataType::string_datatype(), false, false)] + #[case::unsupported_list( + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())), + true, + false + )] + #[case::unsupported_null(ConcreteDataType::null_datatype(), true, false)] + fn test_null_padding_requires_a_nullable_encodable_column( + #[case] data_type: ConcreteDataType, + #[case] nullable: bool, + #[case] known: bool, + ) { + let mut metadata = (*metadata(&[Value::Null])).clone(); + metadata.column_metadatas[0].column_schema = + ColumnSchema::new("tag_0", data_type, nullable); + let metadata = Arc::new( + RegionMetadataBuilder::from_existing(metadata) + .build() + .unwrap(), + ); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + let expected = known.then(|| (Bytes::from_static(&[0]), Bytes::from_static(&[0]))); + assert_eq!(expected, mapper.map((Bytes::new(), Bytes::new())).unwrap()); + if known { + let encoded_null = encode(&metadata, &[Value::Null]); + assert_eq!(Some((encoded_null.clone(), encoded_null)), expected); + } + } + + // Invalid order/layout is an error; foreign bounds and unavailable defaults + // remain unknown. Exercise both endpoints without an inverted range masking decoding. + #[rstest] + #[case::dense(PrimaryKeyEncoding::Dense)] + #[case::sparse(PrimaryKeyEncoding::Sparse)] + fn test_reversed_pk_bounds_return_error(#[case] encoding: PrimaryKeyEncoding) { + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use mito_codec::row_converter::build_primary_key_codec; + + let mut metadata = (*metadata(&[Value::from("")])).clone(); + metadata.primary_key_encoding = encoding; + let metadata = Arc::new(metadata); + let codec = build_primary_key_codec(&metadata); + let mut a = Vec::new(); + let mut b = Vec::new(); + codec + .encode_values(&[(0, Value::from("a"))], &mut a) + .unwrap(); + codec + .encode_values(&[(0, Value::from("b"))], &mut b) + .unwrap(); + let mapper = PrimaryKeyRangeMapper::new(metadata); + let err = mapper.map((b.into(), a.into())).unwrap_err(); + assert!(matches!( + err, + crate::error::Error::InvalidPrimaryKeyRange { .. } + )); + assert_eq!(StatusCode::Internal, err.status_code()); + } + + #[test] + fn test_different_pk_prefix_lengths_return_error() { + let metadata = metadata(&[Value::from(""), Value::Int64(42)]); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + let min = encode(&metadata, &[Value::from("a")]); + let max = encode(&metadata, &[Value::from("b"), Value::Int64(42)]); + assert!(matches!( + mapper.map((min, max)), + Err(crate::error::Error::InvalidPrimaryKeyRange { .. }) + )); + } + + #[rstest] + #[case::min("min")] + #[case::max("max")] + fn test_truncated_pk_endpoint_preserves_decode_error(#[case] endpoint: &'static str) { + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + + let metadata = metadata(&[Value::from("")]); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + let mut min = encode(&metadata, &[Value::from("a")]); + let mut max = encode(&metadata, &[Value::from("b")]); + if endpoint == "min" { + min.truncate(min.len() - 1); + } else { + max.truncate(max.len() - 1); + } + let err = mapper.map((min, max)).unwrap_err(); + assert_eq!(StatusCode::Internal, err.status_code()); + assert!(matches!(err, crate::error::Error::DecodePrimaryKeyRange { + endpoint: actual, + source: mito_codec::error::Error::InvalidDensePrimaryKey { .. }, + .. + } if actual == endpoint)); + } + + #[rstest] + #[case::source_first(false)] + #[case::destination_first(true)] + fn test_aligned_cache_uses_table_schema_across_region_migration( + #[case] destination_first: bool, + ) { + let source = metadata(&[Value::from("")]); + let file = FileHandle::new( + file_meta(&source, &[Value::from("a")]), + new_noop_file_purger(), + ); + let target = metadata_at_version(&[Value::from(""), Value::Int64(42)], 1); + let expected = encode(&target, &[Value::from("a"), Value::Int64(42)]); + let mut regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; + if destination_first { + regions.reverse(); + } + let mut previous: Option<(Bytes, Bytes)> = None; + for region_id in regions { + let mut metadata = (*target).clone(); + metadata.region_id = region_id; + let mapper = PrimaryKeyRangeMapper::new(Arc::new(metadata)); + let aligned = file.primary_key_range(&mapper).unwrap(); + assert_eq!((expected.clone(), expected.clone()), aligned); + if let Some(previous) = previous { + // Different metadata allocations/region ids still share one schema-version cache. + assert_eq!(previous.0.as_ptr(), aligned.0.as_ptr()); + assert_eq!(previous.1.as_ptr(), aligned.1.as_ptr()); + } + previous = Some(aligned); + } + } + + #[test] + fn test_invalid_file_ranges_log_once_and_keep_possible_overlap() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_telemetry::tracing_subscriber::Layer; + use common_telemetry::tracing_subscriber::layer::Context; + use common_telemetry::tracing_subscriber::prelude::*; + use tracing::{Event, Level, Subscriber}; + + struct WarningCounter(Arc); + impl Layer for WarningCounter { + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + if *event.metadata().level() == Level::WARN { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + } + + let warnings = Arc::new(AtomicUsize::new(0)); + let subscriber = + common_telemetry::tracing_subscriber::registry().with(WarningCounter(warnings.clone())); + let metadata = metadata(&[Value::from("")]); + let mapper = Arc::new(PrimaryKeyRangeMapper::new(metadata.clone())); + let healthy_meta = file_meta(&metadata, &[Value::from("c")]); + let healthy = FileHandle::new(healthy_meta, new_noop_file_purger()); + let a = encode(&metadata, &[Value::from("a")]); + let b = encode(&metadata, &[Value::from("b")]); + + tracing::subscriber::with_default(subscriber, || { + for (min, max) in [(b.clone(), a.clone()), (a.slice(..a.len() - 1), b)] { + let mut meta = file_meta(&metadata, &[Value::from("a")]); + meta.primary_key_min = Some(min); + meta.primary_key_max = Some(max); + let file = FileHandle::new(meta, new_noop_file_purger()); + assert_eq!(None, file.primary_key_range(&mapper)); + assert_eq!(None, file.clone().primary_key_range(&mapper)); + // These raw ranges look disjoint from c; unknown must prevent pruning. + assert!(files_overlap_inclusive(&file, &healthy, &mapper)); + assert!(files_overlap_inclusive(&healthy, &file, &mapper)); + } + }); + assert_eq!(2, warnings.load(Ordering::Relaxed)); + } + + #[test] + fn test_missing_impure_default_is_unknown_but_complete_key_is_usable() { + let mut metadata = (*metadata(&[Value::Timestamp(Timestamp::new_millisecond(0))])).clone(); + metadata.column_metadatas[0].column_schema = metadata.column_metadatas[0] + .column_schema + .clone() + .with_default_constraint(Some(ColumnDefaultConstraint::Function( + "current_timestamp()".into(), + ))) + .unwrap(); + let metadata = Arc::new(metadata); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + assert_eq!(None, mapper.map((Bytes::new(), Bytes::new())).unwrap()); + let key = encode( + &metadata, + &[Value::Timestamp(Timestamp::new_millisecond(10))], + ); + assert_eq!( + Some((key.clone(), key.clone())), + mapper.map((key.clone(), key)).unwrap() + ); + } + + #[test] + fn test_sparse_ranges_keep_their_original_encoding() { + use mito_codec::row_converter::SparsePrimaryKeyCodec; + let mut metadata = (*metadata(&[Value::from(""), Value::Null])).clone(); + metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; + let metadata = Arc::new(metadata); + let mut bytes = Vec::new(); + SparsePrimaryKeyCodec::new(&metadata) + .encode_values(&[(0, Value::from("a"))], &mut bytes) + .unwrap(); + let key = Bytes::from(bytes); + let mapper = PrimaryKeyRangeMapper::new(metadata); + assert_eq!( + Some((key.clone(), key.clone())), + mapper.map((key.clone(), key)).unwrap() + ); + } + + #[test] + fn test_aligned_cache_isolates_schemas_and_shares_file_state() { + let old_metadata = metadata(&[Value::from("")]); + let new_metadata = + metadata_at_version(&[Value::from(""), Value::Null, Value::Int64(42)], 1); + let raw_meta = file_meta(&old_metadata, &[Value::from("b")]); + let original = version_with_files(old_metadata.clone(), vec![raw_meta.clone()]); + let old_file = original.ssts.levels()[0].files().next().unwrap(); + let old_mapper = original.ssts.primary_key_mapper(); + let changed = VersionBuilder::from_version(original.clone()) + .metadata(new_metadata.clone()) + .build(); + let new_file = changed.ssts.levels()[0].files().next().unwrap(); + let new_mapper = changed.ssts.primary_key_mapper(); + let expected = encode( + &new_metadata, + &[Value::from("b"), Value::Null, Value::Int64(42)], + ); + assert_eq!( + raw_meta.primary_key_range(), + old_file.primary_key_range(&old_mapper) + ); + assert_eq!( + Some((expected.clone(), expected)), + new_file.primary_key_range(&new_mapper) + ); + assert_eq!(raw_meta, *new_file.meta_ref()); + assert_eq!( + raw_meta.primary_key_range(), + new_file.raw_primary_key_range() + ); + assert_eq!( + old_file.primary_key_range(&old_mapper), + new_file.primary_key_range(&old_mapper) + ); + assert_eq!( + new_file.primary_key_range(&new_mapper), + old_file.primary_key_range(&new_mapper) + ); + old_file.set_compacting(true); + assert!(new_file.compacting()); + new_file.mark_deleted(); + assert!(old_file.is_deleted()); + + // Old-schema flushes can finish after ALTER; compare them with the target mapper. + let added_meta = file_meta(&old_metadata, &[Value::from("a")]); + let added_id = added_meta.file_id; + let changed = VersionBuilder::from_version(Arc::new(changed)) + .apply_edit( + RegionEdit { + files_to_add: vec![added_meta], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + new_noop_file_purger(), + ) + .build(); + let added = &changed.ssts.levels()[0].files[&added_id]; + let expected = encode( + &new_metadata, + &[Value::from("a"), Value::Null, Value::Int64(42)], + ); + assert_eq!( + Some((expected.clone(), expected)), + added.primary_key_range(&new_mapper) + ); + } + + #[test] + fn test_late_statistics_and_changed_defaults_use_each_pinned_schema() { + let metadata_v1 = metadata(&[Value::from(""), Value::Int64(42)]); + let metadata_v2 = metadata_at_version(&[Value::from(""), Value::Int64(100)], 1); + let mut meta = file_meta(&metadata_v1, &[Value::from("a")]); + let raw = meta.primary_key_range().unwrap(); + meta.primary_key_min = None; + meta.primary_key_max = None; + let file = FileHandle::new(meta, new_noop_file_purger()); + let v1 = PrimaryKeyRangeMapper::new(metadata_v1.clone()); + let v2 = PrimaryKeyRangeMapper::new(metadata_v2.clone()); + assert_eq!(None, file.primary_key_range(&v1)); + assert_eq!(None, file.primary_key_range(&v2)); + file.set_primary_key_range(raw.clone()); + for (mapper, metadata, default) in [ + (&v2, &metadata_v2, 100), + (&v1, &metadata_v1, 42), + (&v2, &metadata_v2, 100), + ] { + let expected = encode(metadata, &[Value::from("a"), Value::Int64(default)]); + assert_eq!( + Some((expected.clone(), expected)), + file.primary_key_range(mapper) + ); + } + assert_eq!(Some(raw), file.raw_primary_key_range()); + } + + #[test] + fn test_concurrent_snapshots_do_not_return_each_others_aligned_bounds() { + let old_metadata = metadata(&[Value::from(""), Value::Int64(42)]); + let new_metadata = metadata_at_version(&[Value::from(""), Value::Int64(100)], 1); + let raw = file_meta(&old_metadata, &[Value::from("a")]); + let original_bounds = raw.primary_key_range(); + let file = FileHandle::new(raw, new_noop_file_purger()); + let barrier = std::sync::Barrier::new(2); + std::thread::scope(|scope| { + for (metadata, default) in [(old_metadata, 42), (new_metadata, 100)] { + let file = &file; + let barrier = &barrier; + scope.spawn(move || { + let expected = encode(&metadata, &[Value::from("a"), Value::Int64(default)]); + let mapper = PrimaryKeyRangeMapper::new(metadata); + let mut all_matched = true; + for _ in 0..64 { + barrier.wait(); + all_matched &= file.primary_key_range(&mapper) + == Some((expected.clone(), expected.clone())); + } + // Assert after all barriers so a failure cannot strand the other snapshot. + assert!( + all_matched, + "wrong bounds for schema version {}", + mapper.schema_version() + ); + }); + } + }); + assert_eq!(original_bounds, file.raw_primary_key_range()); + } + + #[test] + fn test_deserialized_compaction_files_use_the_comparison_schema() { + use crate::compaction::CompactionOutput; + use crate::compaction::picker::{PickerOutput, SerializedPickerOutput}; + + let metadata = metadata(&[Value::from(""), Value::Int64(42)]); + let raw = file_meta(&metadata, &[Value::from("a")]); + let picked = PickerOutput { + outputs: vec![CompactionOutput { + output_level: 1, + inputs: vec![FileHandle::new(raw.clone(), new_noop_file_purger())], + filter_deleted: false, + output_time_range: None, + }], + ..Default::default() + }; + let serialized = SerializedPickerOutput::from(&picked); + let output = PickerOutput::from_serialized(serialized, new_noop_file_purger()); + let mapper = PrimaryKeyRangeMapper::new(metadata.clone()); + let file = &output.outputs[0].inputs[0]; + let completed = encode(&metadata, &[Value::from("a"), Value::Int64(42)]); + assert_eq!( + Some((completed.clone(), completed)), + file.primary_key_range(&mapper) + ); + assert_eq!(raw, *file.meta_ref()); + } + + #[test] + fn test_compaction_overlap_uses_completed_endpoints() { + let metadata = metadata(&[Value::from(""), Value::Null]); + let mapper = Arc::new(PrimaryKeyRangeMapper::new(metadata.clone())); + let mut old = file_meta(&metadata, &[Value::from("a")]); + old.primary_key_max = Some(encode(&metadata, &[Value::from("b")])); + let mut new = file_meta(&metadata, &[Value::from("b"), Value::Null]); + new.primary_key_max = Some(encode(&metadata, &[Value::from("c"), Value::Null])); + assert!(old.primary_key_max < new.primary_key_min); + let old = FileHandle::new(old, new_noop_file_purger()); + let new = FileHandle::new(new, new_noop_file_purger()); + assert!(files_overlap_inclusive(&old, &new, &mapper)); + assert!(files_overlap_inclusive(&new, &old, &mapper)); + } +} diff --git a/src/mito2/src/sst/version.rs b/src/mito2/src/sst/version.rs index 172d8701a73..1ca0410db7b 100644 --- a/src/mito2/src/sst/version.rs +++ b/src/mito2/src/sst/version.rs @@ -18,31 +18,45 @@ use std::fmt; use std::sync::Arc; use common_time::{TimeToLive, Timestamp}; +use store_api::metadata::RegionMetadataRef; use store_api::storage::{FileId, RegionId}; use crate::sst::file::{FileHandle, FileMeta, FileTimeRange, Level, MAX_LEVEL}; use crate::sst::file_purger::FilePurgerRef; +use crate::sst::primary_key::PrimaryKeyRangeMapper; /// A version of all SSTs in a region. #[derive(Debug, Clone)] pub(crate) struct SstVersion { /// SST metadata organized by levels. - levels: LevelMetaArray, + levels: Arc, + primary_key_mapper: Arc, } pub(crate) type SstVersionRef = Arc; impl SstVersion { /// Returns a new [SstVersion]. - pub(crate) fn new() -> SstVersion { + pub(crate) fn new(metadata: RegionMetadataRef) -> SstVersion { SstVersion { - levels: new_level_meta_vec(), + levels: Arc::new(new_level_meta_vec()), + primary_key_mapper: Arc::new(PrimaryKeyRangeMapper::new(metadata)), } } + /// Changes the target schema without copying the SST list or rebinding handles. + pub(crate) fn set_metadata(&mut self, metadata: RegionMetadataRef) { + self.primary_key_mapper = Arc::new(self.primary_key_mapper.with_metadata(metadata)); + } + + /// Shares the target schema and encoded defaults with comparisons. + pub(crate) fn primary_key_mapper(&self) -> Arc { + self.primary_key_mapper.clone() + } + /// Returns a slice to metadatas of all levels. pub(crate) fn levels(&self) -> &[LevelMeta] { - &self.levels + self.levels.as_ref() } /// Returns the current handle matching the selected file's identity in its immutable level. @@ -65,11 +79,12 @@ impl SstVersion { file_purger: FilePurgerRef, files_to_add: impl Iterator, ) { + let levels = Arc::make_mut(&mut self.levels); for file in files_to_add { let level = file.level; let new_index_version = file.index_version; // If the file already exists, then we should only replace the handle when the index is outdated. - self.levels[level as usize] + levels[level as usize] .files .entry(file.file_id) .and_modify(|f| { @@ -99,9 +114,10 @@ impl SstVersion { /// # Panics /// Panics if level of [FileMeta] is greater than [MAX_LEVEL]. pub(crate) fn remove_files(&mut self, files_to_remove: impl Iterator) { + let levels = Arc::make_mut(&mut self.levels); for file in files_to_remove { let level = file.level; - if let Some(handle) = self.levels[level as usize].files.remove(&file.file_id) { + if let Some(handle) = levels[level as usize].files.remove(&file.file_id) { handle.mark_deleted(); } } @@ -109,7 +125,7 @@ impl SstVersion { /// Marks all SSTs in this version as deleted. pub(crate) fn mark_all_deleted(&self) { - for level_meta in &self.levels { + for level_meta in self.levels.iter() { for file_handle in level_meta.files.values() { file_handle.mark_deleted(); } @@ -289,7 +305,7 @@ mod tests { ..Default::default() }; - let mut version = SstVersion::new(); + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); version.add_files(purger, [owned, referenced].into_iter()); assert_eq!( @@ -318,7 +334,7 @@ mod tests { ..Default::default() }; - let mut version = SstVersion::new(); + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); version.add_files(purger, [seconds, millis].into_iter()); assert_eq!( @@ -332,7 +348,8 @@ mod tests { #[test] fn time_range_is_none_without_files() { - assert_eq!(SstVersion::new().time_range(), None); + let version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + assert_eq!(version.time_range(), None); } #[test] @@ -346,7 +363,7 @@ mod tests { }) .collect::>(); - let mut version = SstVersion::new(); + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); // files[1] is added multiple times, and that's ok. version.add_files(purger.clone(), files[..=1].iter().cloned()); version.add_files(purger, files[1..].iter().cloned()); @@ -370,7 +387,7 @@ mod tests { }, purger.clone(), ); - let mut version = SstVersion::new(); + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); version.add_files( purger, [ @@ -422,7 +439,7 @@ mod tests { }, ]; - let mut version = SstVersion::new(); + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); version.add_files(purger, files.iter().cloned()); assert_eq!(3, version.owned_num_rows(region_id)); diff --git a/tests-integration/fixtures/docker-compose.yml b/tests-integration/fixtures/docker-compose.yml index 283998cc28f..1e527c27019 100644 --- a/tests-integration/fixtures/docker-compose.yml +++ b/tests-integration/fixtures/docker-compose.yml @@ -134,6 +134,19 @@ services: - MYSQL_USER=greptimedb - MYSQL_PASSWORD=admin - MYSQL_ROOT_PASSWORD=admin + # Wait for authenticated TCP queries, not just the container or initialization socket. + healthcheck: + test: + - CMD-SHELL + - >- + MYSQL_PWD="$${MYSQL_PASSWORD}" mysql + --protocol=TCP --host=127.0.0.1 + --user="$${MYSQL_USER}" --database="$${MYSQL_DATABASE}" + --execute="SELECT 1" + interval: 2s + timeout: 5s + retries: 30 + start_period: 10s volumes: minio_data: From 04f07a91b9ae6baa2442fc112cc2b64e8cbdcc7f Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 21 Sep 2026 02:26:49 +0000 Subject: [PATCH 27/63] ci: use mold in dev-builder (#9262) --- docker/dev-builder/riscv64/Dockerfile | 15 +++++ docker/dev-builder/ubuntu/Dockerfile | 13 +++++ docker/dev-builder/ubuntu/Dockerfile-20.04 | 66 ---------------------- 3 files changed, 28 insertions(+), 66 deletions(-) delete mode 100644 docker/dev-builder/ubuntu/Dockerfile-20.04 diff --git a/docker/dev-builder/riscv64/Dockerfile b/docker/dev-builder/riscv64/Dockerfile index 8505fa3b49a..6ea5c5d68ee 100644 --- a/docker/dev-builder/riscv64/Dockerfile +++ b/docker/dev-builder/riscv64/Dockerfile @@ -11,6 +11,7 @@ RUN apt-get update && \ # Install dependencies, including the riscv64 cross toolchain. RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ libssl-dev \ + libatomic1 \ tzdata \ curl \ unzip \ @@ -40,6 +41,20 @@ RUN mv -f protoc3/bin/* /usr/local/bin/ # contains protobuf headers under /usr/local/include. RUN cp -r protoc3/include/* /usr/local/include/ +# Install mold and make it the default linker (the equivalent of +# rui314/setup-mold with its default options on GitHub-hosted runners) for +# lower peak memory usage during linking. The host `ld` symlink covers +# host-side linking (build scripts, proc macros); the riscv64 symlink covers +# the cross target -- the cross gcc driver passes -m elf64lriscv explicitly, +# so mold picks the right emulation. +ARG MOLD_VERSION=2.42.1 +RUN arch="$(uname -m)" && \ + curl -fsSL -o /tmp/mold.tar.gz https://github.com/rui314/mold/releases/download/v${MOLD_VERSION}/mold-${MOLD_VERSION}-${arch}-linux.tar.gz && \ + tar -C /usr/local --strip-components=1 -xzf /tmp/mold.tar.gz && \ + rm /tmp/mold.tar.gz && \ + ln -sf /usr/local/bin/mold "$(readlink -f /usr/bin/ld)" && \ + ln -sf /usr/local/bin/mold "$(readlink -f "$(command -v riscv64-linux-gnu-ld)")" + # Silence all `safe.directory` warnings, to avoid the "detect dubious repository" # error when building with the repository mounted from a different user. RUN git config --global --add safe.directory '*' diff --git a/docker/dev-builder/ubuntu/Dockerfile b/docker/dev-builder/ubuntu/Dockerfile index 198797b5326..ac3e2954393 100644 --- a/docker/dev-builder/ubuntu/Dockerfile +++ b/docker/dev-builder/ubuntu/Dockerfile @@ -11,6 +11,7 @@ RUN apt-get update && \ # Install dependencies. RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ libssl-dev \ + libatomic1 \ tzdata \ curl \ unzip \ @@ -35,6 +36,18 @@ fi RUN mv protoc3/bin/* /usr/local/bin/ RUN mv protoc3/include/* /usr/local/include/ +# Install mold and make it the default linker (the equivalent of +# rui314/setup-mold with its default options on GitHub-hosted runners), so +# every link inside this image -- binaries, test executables, build scripts, +# proc macros -- goes through mold for lower peak memory usage. The prebuilt +# binaries require GLIBC 2.18+. +ARG MOLD_VERSION=2.42.1 +RUN arch="$(uname -m)" && \ + curl -fsSL -o /tmp/mold.tar.gz https://github.com/rui314/mold/releases/download/v${MOLD_VERSION}/mold-${MOLD_VERSION}-${arch}-linux.tar.gz && \ + tar -C /usr/local --strip-components=1 -xzf /tmp/mold.tar.gz && \ + rm /tmp/mold.tar.gz && \ + ln -sf /usr/local/bin/mold "$(readlink -f /usr/bin/ld)" + # Silence all `safe.directory` warnings, to avoid the "detect dubious repository" error when building with submodules. # Disabling the safe directory check here won't pose extra security issues, because in our usage for this dev build # image, we use it solely on our own environment (that github action's VM, or ECS created dynamically by ourselves), diff --git a/docker/dev-builder/ubuntu/Dockerfile-20.04 b/docker/dev-builder/ubuntu/Dockerfile-20.04 deleted file mode 100644 index d78046698c3..00000000000 --- a/docker/dev-builder/ubuntu/Dockerfile-20.04 +++ /dev/null @@ -1,66 +0,0 @@ -FROM ubuntu:20.04 - -# The root path under which contains all the dependencies to build this Dockerfile. -ARG DOCKER_BUILD_ROOT=. - -ENV LANG en_US.utf8 -WORKDIR /greptimedb - -RUN apt-get update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y software-properties-common -# Install dependencies. -RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ - libssl-dev \ - tzdata \ - curl \ - unzip \ - ca-certificates \ - git \ - build-essential \ - pkg-config - -ARG TARGETPLATFORM -RUN echo "target platform: $TARGETPLATFORM" - -ARG PROTOBUF_VERSION=29.3 - -# Install protobuf, because the one in the apt is too old (v3.12). -RUN if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \ - curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip && \ - unzip protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip -d protoc3; \ -elif [ "$TARGETPLATFORM" = "linux/amd64" ]; then \ - curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-x86_64.zip && \ - unzip protoc-${PROTOBUF_VERSION}-linux-x86_64.zip -d protoc3; \ -fi -RUN mv protoc3/bin/* /usr/local/bin/ -RUN mv protoc3/include/* /usr/local/include/ - -# Silence all `safe.directory` warnings, to avoid the "detect dubious repository" error when building with submodules. -# Disabling the safe directory check here won't pose extra security issues, because in our usage for this dev build -# image, we use it solely on our own environment (that github action's VM, or ECS created dynamically by ourselves), -# and the repositories are pulled from trusted sources (still us, of course). Doing so does not violate the intention -# of the Git's addition to the "safe.directory" at the first place (see the commit message here: -# https://github.com/git/git/commit/8959555cee7ec045958f9b6dd62e541affb7e7d9). -# There's also another solution to this, that we add the desired submodules to the safe directory, instead of using -# wildcard here. However, that requires the git's config files and the submodules all owned by the very same user. -# It's troublesome to do this since the dev build runs in Docker, which is under user "root"; while outside the Docker, -# it can be a different user that have prepared the submodules. -RUN git config --global --add safe.directory '*' - -# Install Rust. -SHELL ["/bin/bash", "-c"] -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --no-modify-path --default-toolchain none -y -ENV PATH /root/.cargo/bin/:$PATH - -# Install Rust toolchains. -ARG RUST_TOOLCHAIN -RUN rustup toolchain install ${RUST_TOOLCHAIN} - -# Install cargo-binstall with a specific version to adapt the current rust toolchain. -# Note: if we use the latest version, we may encounter the following `use of unstable library feature 'io_error_downcast'` error. -# compile from source take too long, so we use the precompiled binary instead -COPY $DOCKER_BUILD_ROOT/docker/dev-builder/binstall/pull_binstall.sh /usr/local/bin/pull_binstall.sh -RUN chmod +x /usr/local/bin/pull_binstall.sh && /usr/local/bin/pull_binstall.sh - -# Install nextest. -RUN cargo binstall cargo-nextest --no-confirm From 055bbbc93d02a78ff6154eacc3ae02042ef86c31 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 21 Sep 2026 02:35:05 +0000 Subject: [PATCH 28/63] feat: update pgwire to 0.41 (#9025) * feat: update pgwire to 0.41 * feat: update datafusion-pg-catalog and arrow-pg --- Cargo.lock | 12 ++++++------ Cargo.toml | 4 ++-- src/servers/Cargo.toml | 4 ++-- src/servers/src/postgres/handler.rs | 19 ++++++++----------- 4 files changed, 18 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b2a7b3e1035..2abfa9b7c37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -653,7 +653,7 @@ dependencies = [ [[package]] name = "arrow-pg" version = "0.15.0" -source = "git+https://github.com/GreptimeTeam/datafusion-postgres.git?rev=3c77e6c32b8db80635a0d2f4b318a36b31170bc3#3c77e6c32b8db80635a0d2f4b318a36b31170bc3" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4280,7 +4280,7 @@ dependencies = [ [[package]] name = "datafusion-pg-catalog" version = "0.18.3" -source = "git+https://github.com/GreptimeTeam/datafusion-postgres.git?rev=3c77e6c32b8db80635a0d2f4b318a36b31170bc3#3c77e6c32b8db80635a0d2f4b318a36b31170bc3" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" dependencies = [ "arrow-pg", "async-trait", @@ -10698,9 +10698,9 @@ dependencies = [ [[package]] name = "pg_interval" -version = "0.5.0" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c386dd54fce258fc04e668126ae68589a0d92e03a90ea67881d1300f70fd6170" +checksum = "8408fb1e08ebf2e7676d001caa6482d1b33e7e1f1dad9d3a758a05640eea2026" dependencies = [ "bytes", "chrono", @@ -10709,9 +10709,9 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.40.7" +version = "0.41.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8265901ede50d0879fe401c6fe282e7e4ff83ae10a48c1f8e4f89b92f7d5f604" +checksum = "c667f8f4046889d52ed11e17aae954b1816d3df48dd23bdebcf199cef332df93" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/Cargo.toml b/Cargo.toml index f1b5577d391..54fddcce397 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -376,8 +376,8 @@ datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.g datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-pg-catalog = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" } -arrow-pg = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" } +datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } +arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } # Temporary: use the GreptimeTeam fork of tikv-jemalloc-sys embedding diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 3baba2d6634..5fa10313e04 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -95,8 +95,8 @@ operator.workspace = true otel-arrow-rust.workspace = true parking_lot.workspace = true partition.workspace = true -pg_interval = "0.5" -pgwire = { version = "0.40", default-features = false, features = [ +pg_interval = "0.6" +pgwire = { version = "0.41", default-features = false, features = [ "server-api-aws-lc-rs", "pg-ext-types", ] } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index aca4362bd34..204fa7745a5 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -317,23 +317,20 @@ impl QueryParser for DefaultQueryParser { _client: &C, sql: &str, _types: &[Option], - ) -> PgWireResult { + ) -> PgWireResult> { crate::metrics::METRIC_POSTGRES_PREPARED_COUNT.inc(); let query_ctx = self.session.new_query_context(); // do not parse if query is empty or matches rules if sql.is_empty() { - return Ok(PgSqlPlan { - plan: SqlPlan::Empty, - copy_to_stdout_format: None, - }); + return Ok(None); } if fixtures::matches(sql) { - return Ok(PgSqlPlan { + return Ok(Some(PgSqlPlan { plan: SqlPlan::Shortcut(sql.to_string()), copy_to_stdout_format: None, - }); + })); } let parsed_statements = self.compatibility_parser.parse(sql); @@ -367,15 +364,15 @@ impl QueryParser for DefaultQueryParser { .map_err(convert_err)? .map(|DescribeResult { logical_plan }| logical_plan) { - Ok(PgSqlPlan { + Ok(Some(PgSqlPlan { plan: SqlPlan::Plan(logical_plan, stmt), copy_to_stdout_format, - }) + })) } else { - Ok(PgSqlPlan { + Ok(Some(PgSqlPlan { plan: SqlPlan::Statement(stmt, sql), copy_to_stdout_format, - }) + })) } } } From 941193e9c169ab589d88b8e3e8bc1e365cd686c2 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 21 Sep 2026 03:37:47 +0000 Subject: [PATCH 29/63] refactor: reorganize logical table batching and isolate encoding (#9210) * refactor: relocate the logical table batcher Signed-off-by: WenyXu * refactor: isolate logical batch conversion and region writes Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/frontend/src/instance/prom_store.rs | 2 +- src/frontend/src/server.rs | 8 +- src/servers/benches/flush_batch_physical.rs | 4 +- src/servers/src/batcher.rs | 3 +- .../logical_table.rs} | 1496 ++--------------- .../batcher/logical_table/batch_convert.rs | 686 ++++++++ .../src/batcher/logical_table/region_write.rs | 599 +++++++ .../src/batcher/logical_table/test_util.rs | 187 +++ src/servers/src/http.rs | 4 +- src/servers/src/http/prom_store.rs | 12 +- src/servers/src/lib.rs | 1 - src/servers/src/prom_row_builder.rs | 2 +- tests-integration/src/test_util.rs | 4 +- 13 files changed, 1607 insertions(+), 1401 deletions(-) rename src/servers/src/{pending_rows_batcher.rs => batcher/logical_table.rs} (65%) create mode 100644 src/servers/src/batcher/logical_table/batch_convert.rs create mode 100644 src/servers/src/batcher/logical_table/region_write.rs create mode 100644 src/servers/src/batcher/logical_table/test_util.rs diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index e0b17b1a89e..9aa228715c0 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/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, diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 257e5243e89..532b0d08b11 100644 --- a/src/frontend/src/server.rs +++ b/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); diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index 45a3f37c780..be992692d13 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/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; diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs index c87ebd2eff4..10f8f558eb3 100644 --- a/src/servers/src/batcher.rs +++ b/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)] diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/batcher/logical_table.rs similarity index 65% rename from src/servers/src/pending_rows_batcher.rs rename to src/servers/src/batcher/logical_table.rs index 79ae645b98b..de872ac326b 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -12,20 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod batch_convert; +mod region_write; +#[cfg(test)] +mod test_util; + use std::collections::{HashMap, HashSet}; use std::num::NonZeroUsize; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; -use api::v1::meta::Peer; -use api::v1::region::{ - BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, -}; -use api::v1::{ArrowIpc, ColumnSchema, RowInsertRequests, Rows}; -use arrow::compute::{concat_batches, filter_record_batch}; -use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; -use arrow::record_batch::RecordBatch; +use api::v1::{ColumnSchema, RowInsertRequests, Rows}; +use arrow::datatypes::Schema as ArrowSchema; use async_trait::async_trait; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; @@ -35,25 +34,31 @@ use common_batcher::notifier::{Notifier, run_notifier}; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; -use common_grpc::error::Error as GrpcError; -use common_grpc::flight::record_batch_to_ipc; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; -use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_timestamp, greptime_value}; +use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_runtime::spawn_global; -use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; use datatypes::timestamp::append_timestamps; -use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; use partition::manager::PartitionRuleManagerRef; -use partition::partition::PartitionRuleRef; use session::context::QueryContextRef; -use smallvec::SmallVec; -use snafu::{OptionExt, ResultExt, ensure}; -use store_api::storage::{RegionId, TableId}; -use table::metadata::{TableInfo, TableInfoRef}; +use snafu::{OptionExt, ResultExt}; +use table::metadata::TableId; use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, mpsc, oneshot}; +pub use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; +use crate::batcher::logical_table::batch_convert::{ + concat_modified_batches, transform_logical_batches_to_physical, +}; +use crate::batcher::logical_table::region_write::{ + CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter, + PartitionManagerPhysicalFlushAdapter, encode_region_write_requests, + flush_region_writes_concurrently, plan_region_batches, resolve_region_targets, +}; +pub use crate::batcher::logical_table::region_write::{ + PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalTableMetadata, +}; use crate::error; use crate::error::{Error, Result}; use crate::metrics::{ @@ -67,6 +72,7 @@ use crate::prom_row_builder::{ }; const PHYSICAL_TABLE_KEY: &str = "physical_table"; + /// Whether wait for ingestion result before reply to client. const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; @@ -84,12 +90,14 @@ pub fn pending_rows_batch_sync_enabled() -> bool { .and_then(|v| v.parse::().ok()) .unwrap_or(true) } + const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; -const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3; + const MAX_CONCURRENT_FLOW_NOTIFICATIONS: NonZeroUsize = NonZeroUsize::new(8).unwrap(); + #[async_trait] pub trait PendingRowsSchemaAlterer: Send + Sync { - /// Batch-create multiple logical tables that are missing. + /// FlushBatch-create multiple logical tables that are missing. /// Each entry is `(table_name, request_schema)`. async fn create_tables_if_missing_batch( &self, @@ -100,7 +108,7 @@ pub trait PendingRowsSchemaAlterer: Send + Sync { ctx: QueryContextRef, ) -> Result<()>; - /// Batch-alter multiple logical tables to add missing tag columns. + /// FlushBatch-alter multiple logical tables to add missing tag columns. /// Each entry is `(table_name, missing_column_names)`. async fn add_missing_prom_tag_columns_batch( &self, @@ -113,116 +121,6 @@ pub trait PendingRowsSchemaAlterer: Send + Sync { pub type PendingRowsSchemaAltererRef = Arc; -#[derive(Clone)] -pub struct PhysicalTableMetadata { - pub table_info: TableInfoRef, - /// Mapping from column name to column id - pub col_name_to_ids: Option>, -} - -#[async_trait] -pub trait PhysicalFlushCatalogProvider: Send + Sync { - async fn physical_table( - &self, - catalog: &str, - schema: &str, - table_name: &str, - query_ctx: &session::context::QueryContext, - ) -> catalog::error::Result>; -} - -#[async_trait] -pub trait PhysicalFlushPartitionProvider: Send + Sync { - async fn find_table_partition_rule( - &self, - table_info: &TableInfo, - ) -> partition::error::Result; - - async fn find_region_leader(&self, region_id: RegionId) -> Result; -} - -#[async_trait] -pub trait PhysicalFlushNodeRequester: Send + Sync { - async fn handle( - &self, - peer: &Peer, - request: RegionRequest, - ) -> Result; -} - -#[derive(Clone)] -struct CatalogManagerPhysicalFlushAdapter { - catalog_manager: CatalogManagerRef, -} - -#[async_trait] -impl PhysicalFlushCatalogProvider for CatalogManagerPhysicalFlushAdapter { - async fn physical_table( - &self, - catalog: &str, - schema: &str, - table_name: &str, - query_ctx: &session::context::QueryContext, - ) -> catalog::error::Result> { - self.catalog_manager - .table(catalog, schema, table_name, Some(query_ctx)) - .await - .map(|table| { - table.map(|table| { - let table_info = table.table_info(); - let name_to_ids = table_info.name_to_ids(); - PhysicalTableMetadata { - table_info, - col_name_to_ids: name_to_ids, - } - }) - }) - } -} - -#[derive(Clone)] -struct PartitionManagerPhysicalFlushAdapter { - partition_manager: PartitionRuleManagerRef, -} - -#[async_trait] -impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter { - async fn find_table_partition_rule( - &self, - table_info: &TableInfo, - ) -> partition::error::Result { - 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) - } -} - -#[derive(Clone)] -struct NodeManagerPhysicalFlushAdapter { - node_manager: NodeManagerRef, -} - -#[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) - .await - .context(error::CommonMetaSnafu) - } -} - #[derive(Debug, Clone, Hash, Eq, PartialEq)] struct BatchKey { catalog: String, @@ -231,59 +129,6 @@ struct BatchKey { skip_wal: bool, } -/// An aligned logical record batch and its timestamp column index. -#[derive(Debug, Clone)] -pub struct RecordBatchWithTsIdx { - /// The aligned logical record batch. - batch: RecordBatch, - /// The timestamp column index in `batch`. - timestamp_index: usize, -} - -impl RecordBatchWithTsIdx { - /// Creates a record batch with a validated timestamp column index. - pub fn try_new(batch: RecordBatch, timestamp_index: usize) -> Result { - let schema = batch.schema(); - let timestamp_field = schema.fields().get(timestamp_index).with_context(|| { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Timestamp column index {} is out of bounds for record batch with {} columns", - timestamp_index, - batch.num_columns() - ), - } - })?; - ensure!( - matches!(timestamp_field.data_type(), ArrowDataType::Timestamp(_, _)), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Column at index {} is not a timestamp column: {:?}", - timestamp_index, - timestamp_field.data_type() - ), - } - ); - - Ok(Self { - batch, - timestamp_index, - }) - } - - #[cfg(test)] - pub(crate) fn into_parts(self) -> (RecordBatch, usize) { - (self.batch, self.timestamp_index) - } -} - -#[derive(Debug, Clone)] -pub struct TableBatch { - pub table_name: String, - pub table_id: TableId, - pub batches: Vec, - pub row_count: usize, -} - /// Intermediate planning state for resolving and preparing logical tables /// before row-to-batch alignment. struct TableResolutionPlan { @@ -346,7 +191,7 @@ fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { } /// Prometheus remote write pending rows batcher. -pub struct PendingRowsBatcher { +pub struct LogicalTablePendingRowsBatcher { workers: Arc>, flush_interval: Duration, flush_policy: TimingFlushPolicy, @@ -363,7 +208,7 @@ pub struct PendingRowsBatcher { shutdown: broadcast::Sender<()>, } -impl PendingRowsBatcher { +impl LogicalTablePendingRowsBatcher { #[allow(clippy::too_many_arguments)] pub fn try_new( partition_manager: PartitionRuleManagerRef, @@ -417,7 +262,9 @@ impl PendingRowsBatcher { shutdown, })) } +} +impl LogicalTablePendingRowsBatcher { pub async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { let (table_batches, total_rows) = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED @@ -496,7 +343,9 @@ impl PendingRowsBatcher { Ok(total_rows as u64) } } +} +impl LogicalTablePendingRowsBatcher { /// Converts proto `RowInsertRequests` directly into aligned `RecordBatch`es /// in a single pass, handling table creation, schema alteration, column /// renaming, reordering, and null-filling without building intermediate @@ -535,7 +384,9 @@ impl PendingRowsBatcher { Ok((aligned_batches, total_rows)) } +} +impl LogicalTablePendingRowsBatcher { /// Extracts non-empty `(table_name, rows)` pairs and computes total row /// count across the retained entries. fn collect_non_empty_table_rows(requests: RowInsertRequests) -> (Vec<(String, Rows)>, usize) { @@ -556,7 +407,9 @@ impl PendingRowsBatcher { (table_rows, total_rows) } +} +impl LogicalTablePendingRowsBatcher { /// Returns unique `(table_name, proto_schema)` pairs while keeping the /// first-seen schema for duplicate table names. fn collect_unique_table_schemas( @@ -582,7 +435,9 @@ impl PendingRowsBatcher { Ok(unique_tables) } +} +impl LogicalTablePendingRowsBatcher { /// Resolves table metadata and classifies each table into existing, /// to-create, and to-alter groups used by subsequent DDL steps. async fn plan_table_resolution( @@ -643,8 +498,10 @@ impl PendingRowsBatcher { Ok(plan) } +} - /// Batch-creates missing tables, refreshes their schema metadata, and +impl LogicalTablePendingRowsBatcher { + /// FlushBatch-creates missing tables, refreshes their schema metadata, and /// enqueues follow-up alters for extra tag columns discovered in later rows. async fn create_missing_tables_and_refresh_schemas( &self, @@ -710,7 +567,9 @@ impl PendingRowsBatcher { Ok(()) } +} +impl LogicalTablePendingRowsBatcher { /// For newly created tables, re-checks all row schemas and appends alter /// operations when additional tag columns are still missing. fn enqueue_alter_for_new_tables( @@ -748,8 +607,10 @@ impl PendingRowsBatcher { Ok(()) } +} - /// Batch-alters tables that have missing tag columns and refreshes the +impl LogicalTablePendingRowsBatcher { + /// FlushBatch-alters tables that have missing tag columns and refreshes the /// in-memory schema map used for row alignment. async fn alter_tables_and_refresh_schemas( &self, @@ -805,7 +666,9 @@ impl PendingRowsBatcher { Ok(()) } +} +impl LogicalTablePendingRowsBatcher { /// Converts proto rows to `RecordBatch` values aligned to resolved region /// schemas and returns `(table_name, table_id, batch)` tuples. fn build_aligned_batches( @@ -832,6 +695,9 @@ impl PendingRowsBatcher { Ok(aligned_batches) } +} + +impl LogicalTablePendingRowsBatcher { async fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker { let (tx, receiver) = self .workers @@ -843,7 +709,9 @@ impl PendingRowsBatcher { } PendingWorker { tx } } +} +impl LogicalTablePendingRowsBatcher { fn spawn_worker( &self, key: BatchKey, @@ -872,7 +740,7 @@ impl PendingRowsBatcher { } } -impl Drop for PendingRowsBatcher { +impl Drop for LogicalTablePendingRowsBatcher { fn drop(&mut self) { let _ = self.shutdown.send(()); } @@ -1116,186 +984,6 @@ async fn spawn_flush( } } -struct FlushRegionWrite { - datanode: Peer, - request: RegionRequest, -} - -struct PlannedRegionBatch { - region_id: RegionId, - batch: RecordBatch, -} - -#[cfg(test)] -impl PlannedRegionBatch { - fn num_rows(&self) -> usize { - self.batch.num_rows() - } -} - -struct ResolvedRegionBatch { - planned: PlannedRegionBatch, - datanode: Peer, -} - -fn should_dispatch_concurrently(region_write_count: usize) -> bool { - region_write_count > 1 -} - -/// Classifies columns in a logical-table batch for sparse primary-key conversion. -/// -/// Returns: -/// - `Vec`: all Utf8 tag columns sorted by tag name, used for -/// TSID and sparse primary-key encoding. -/// - `SmallVec<[usize; 3]>`: indices of columns copied into the physical batch -/// after `__primary_key`, ordered as `[greptime_timestamp, greptime_value, -/// partition_tag_columns...]`. -fn columns_taxonomy( - batch_schema: &Arc, - table_name: &str, - name_to_ids: &HashMap, - partition_columns: &HashSet<&str>, -) -> Result<(Vec, SmallVec<[usize; 3]>)> { - let mut tag_columns = Vec::new(); - let mut essential_column_indices = - SmallVec::<[usize; 3]>::with_capacity(2 + partition_columns.len()); - // Placeholder for greptime_timestamp and greptime_value - essential_column_indices.push(0); - essential_column_indices.push(0); - - let mut timestamp_index = None; - let mut value_index = None; - - for (index, field) in batch_schema.fields().iter().enumerate() { - match field.data_type() { - ArrowDataType::Utf8 => { - let column_id = name_to_ids.get(field.name()).copied().with_context(|| { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Column '{}' from logical table '{}' not found in physical table column IDs", - field.name(), - table_name - ), - } - })?; - tag_columns.push(TagColumnInfo { - name: field.name().clone(), - index, - column_id, - }); - - if partition_columns.contains(field.name().as_str()) { - essential_column_indices.push(index); - } - } - ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => { - ensure!( - timestamp_index.replace(index).is_none(), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Duplicated timestamp column in logical table '{}' batch schema", - table_name - ), - } - ); - } - ArrowDataType::Float64 => { - ensure!( - value_index.replace(index).is_none(), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Duplicated value column in logical table '{}' batch schema", - table_name - ), - } - ); - } - datatype => { - return error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Unexpected data type '{datatype:?}' in logical table '{}' batch schema", - table_name - ), - } - .fail(); - } - } - } - - let timestamp_index = - timestamp_index.with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Missing essential column '{}' in logical table '{}' batch schema", - greptime_timestamp(), - table_name - ), - })?; - let value_index = value_index.with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Missing essential column '{}' in logical table '{}' batch schema", - greptime_value(), - table_name - ), - })?; - - tag_columns.sort_by(|a, b| a.name.cmp(&b.name)); - - essential_column_indices[0] = timestamp_index; - essential_column_indices[1] = value_index; - - Ok((tag_columns, essential_column_indices)) -} - -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) -} - -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) -} - async fn flush_batch_with_managers( flush: FlushBatch, partition_manager: PartitionRuleManagerRef, @@ -1595,220 +1283,6 @@ pub async fn flush_batch_physical( flush_region_writes_concurrently(node_manager, region_writes).await } -/// Transforms logical table batches into physical format (sparse primary key encoding). -/// -/// It identifies tag columns and essential columns (timestamp, value) for each logical batch -/// and applies sparse primary key modification. -fn transform_logical_batches_to_physical( - table_batches: &[TableBatch], - name_to_ids: &HashMap, - partition_columns_set: &HashSet<&str>, -) -> Result> { - let mut modified_batches: Vec = - Vec::with_capacity(table_batches.iter().map(|b| b.batches.len()).sum()); - - let mut modify_elapsed = Duration::ZERO; - let mut columns_taxonomy_elapsed = Duration::ZERO; - - for table_batch in table_batches { - let table_id = table_batch.table_id; - - for batch in &table_batch.batches { - let batch = &batch.batch; - let batch_schema = batch.schema(); - let start = Instant::now(); - let (tag_columns, essential_col_indices) = columns_taxonomy( - &batch_schema, - &table_batch.table_name, - name_to_ids, - partition_columns_set, - )?; - - columns_taxonomy_elapsed += start.elapsed(); - if tag_columns.is_empty() && essential_col_indices.is_empty() { - continue; - } - - let modified = { - let start = Instant::now(); - // The schema of modified batch is: __primary_key, timestamp, value, other partition columns... - let batch = modify_batch_sparse( - batch.clone(), - table_id, - &tag_columns, - &essential_col_indices, - )?; - modify_elapsed += start.elapsed(); - batch - }; - - modified_batches.push(modified); - } - } - - PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_modify_batch"]) - .observe(modify_elapsed.as_secs_f64()); - PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_columns_taxonomy"]) - .observe(columns_taxonomy_elapsed.as_secs_f64()); - - ensure!( - !modified_batches.is_empty(), - error::InternalSnafu { - err_msg: "No batches can be transformed during pending flush", - } - ); - Ok(modified_batches) -} - -/// Concatenates all modified batches into a single large batch. -/// -/// All modified batches share the same physical schema. -fn concat_modified_batches(modified_batches: &[RecordBatch]) -> Result { - let combined_schema = modified_batches[0].schema(); - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_concat_all"]) - .start_timer(); - concat_batches(&combined_schema, modified_batches).context(error::ArrowSnafu) -} - -fn split_combined_batch_by_region( - combined_batch: &RecordBatch, - partition_rule: &dyn partition::partition::PartitionRule, -) -> Result> { - let _timer = 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) -} - -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) -} - -fn plan_region_batch( - stripped_batch: &RecordBatch, - physical_table_id: TableId, - region_number: u32, - mask: &partition::partition::RegionMask, -) -> Result> { - if mask.select_none() { - return Ok(None); - } - - let region_batch = if mask.select_all() { - stripped_batch.clone() - } else { - let _timer = 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)? - }; - - let row_count = region_batch.num_rows(); - if row_count == 0 { - return Ok(None); - } - - Ok(Some(PlannedRegionBatch { - region_id: RegionId::new(physical_table_id, region_number), - batch: region_batch, - })) -} - -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)?; - - 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)? - { - planned_batches.push(planned_batch); - } - } - - Ok(planned_batches) -} - -async fn resolve_region_targets( - 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) -} - -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; - let (schema_bytes, data_header, payload) = { - let _timer = 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)? - }; - - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal, - region_id: region_id.as_u64(), - partition_expr_version: None, - // 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, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema: schema_bytes, - data_header, - payload, - })), - })), - }; - - region_writes.push(FlushRegionWrite { - datanode: resolved.datanode, - request, - }); - } - - Ok(region_writes) -} - fn notify_waiters(waiters: Vec, result: Result<()>) { let shared_result = result.map_err(Arc::new); for waiter in waiters { @@ -1820,41 +1294,37 @@ fn notify_waiters(waiters: Vec, result: Result<()>) { } } -fn map_ipc_error(error: GrpcError) -> Error { - match error { - GrpcError::NotSupported { feat } => Error::NotSupported { feat }, - GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason }, - error => Error::Internal { - err_msg: error.to_string(), - }, - } -} - #[cfg(test)] mod tests { use std::any::Any; - use std::collections::{HashMap, HashSet}; - use std::future::{Future, poll_fn}; + use std::collections::HashMap; + use std::future::poll_fn; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::task::Poll; use std::time::Duration; use api::region::RegionResponse; - use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse}; use api::v1::meta::Peer; - use api::v1::region::{InsertRequests, RegionRequest, region_request}; + use api::v1::region::RegionRequest; use api::v1::value::ValueData; use api::v1::{ ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, Value, }; - use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray}; + use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; use catalog::error::Result as CatalogResult; use catalog::memory::MemoryCatalogManager; + use common_batcher::flush_limiter::FlushLimiter; + use common_batcher::flush_policy::FlushTrigger; + use common_batcher::flush_policy::timing::TimingFlushPolicy; + use common_batcher::notifier::Notifier; + use common_batcher::pending_worker::PendingWorker as PendingCore; + use common_batcher::request_limiter::RequestLimiter; + use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::{ TableFlownodeSetCacheRef, new_table_flownode_set_cache, new_table_route_cache, }; @@ -1862,47 +1332,44 @@ mod tests { use common_meta::instruction::{CacheIdent, CreateFlow}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackend, TxnService}; - use common_meta::node_manager::{ - Datanode, DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, - NodeManagerRef, - }; + use common_meta::node_manager::NodeManagerRef; use common_meta::rpc::store::{ BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse, }; - use common_query::request::QueryRequest; - use common_recordbatch::SendableRecordBatchStream; + use common_query::prelude::greptime_timestamp; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use moka::future::CacheBuilder; use partition::cache::new_partition_info_cache; use partition::error::Result as PartitionResult; use partition::manager::PartitionRuleManager; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; - use smallvec::SmallVec; - use snafu::ResultExt; use store_api::storage::RegionId; use table::metadata::TableId; use table::test_util::table_info::test_table_info; use tokio::sync::{Notify, Semaphore, broadcast, mpsc, oneshot}; - use tokio::time::{advance, sleep}; + use tokio::time::advance; - use crate::error; - use crate::metrics::FLOW_NOTIFICATION_DROPPED; - use crate::pending_rows_batcher::{ - 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, - 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, + use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; + use crate::batcher::logical_table::region_write::{ + PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalTableMetadata, + }; + use crate::batcher::logical_table::test_util::{ + ConcurrentMockNodeManager, FlowNotificationMockNodeManager, RecordingFlownode, + mock_aligned_tag_batch, mock_table_flownode_cache, mock_tag_batch, + }; + use crate::batcher::logical_table::{ + BatchKey, FlushBatch, FlushWaiter, LogicalTablePendingRowsBatcher, PendingBatch, + WorkerCommand, drain_batch, extract_timestamps, flush_batch, flush_batch_physical, + notify_flow_dirty_windows_after_flush, notify_waiters, remove_worker_if_same_channel, + should_close_worker_on_idle_timeout, start_flow_notification_worker, start_worker, try_enqueue_flow_notification, }; + use crate::error; + use crate::error::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 { @@ -1915,37 +1382,6 @@ mod tests { } } - fn mock_tag_batch(tag_name: &str, tag_value: &str, ts: i64, val: f64) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new(tag_name, ArrowDataType::Utf8, true), - ])); - - RecordBatch::try_new( - schema, - vec![ - Arc::new(TimestampMillisecondArray::from(vec![ts])), - Arc::new(arrow::array::Float64Array::from(vec![val])), - Arc::new(StringArray::from(vec![tag_value])), - ], - ) - .unwrap() - } - - fn mock_aligned_tag_batch( - tag_name: &str, - tag_value: &str, - ts: i64, - val: f64, - ) -> RecordBatchWithTsIdx { - RecordBatchWithTsIdx::try_new(mock_tag_batch(tag_name, tag_value, ts, val), 0).unwrap() - } - fn mock_timestamp_batch(timestamps: Vec>) -> RecordBatchWithTsIdx { let batch = RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![Field::new( @@ -2127,7 +1563,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], }; @@ -2232,42 +1668,6 @@ mod tests { } } - struct TwoRegionPartitionRule { - partition_columns: Vec, - } - - impl PartitionRule for TwoRegionPartitionRule { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn partition_columns(&self) -> &[String] { - &self.partition_columns - } - - fn find_region( - &self, - _values: &[datatypes::prelude::Value], - ) -> partition::error::Result { - unimplemented!() - } - - fn split_record_batch( - &self, - _record_batch: &RecordBatch, - ) -> partition::error::Result> - { - Ok(HashMap::from([ - (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), - (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), - ( - 3, - RegionMask::new(BooleanArray::from(vec![false, false]), 0), - ), - ])) - } - } - struct MockFlushPartitionProvider { partition_rule_calls: Arc, region_leader_calls: Arc, @@ -2334,7 +1734,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()); @@ -2460,148 +1861,10 @@ mod tests { assert_eq!("cpu", pending_batch.tables[&43].table_name); } - #[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!() - } - } - - #[derive(Clone)] - struct ConcurrentMockNodeManager { - datanodes: Arc>, - } - - #[async_trait] - impl DatanodeManager for ConcurrentMockNodeManager { - async fn datanode(&self, node: &Peer) -> DatanodeRef { - self.datanodes - .get(&node.id) - .expect("datanode not found") - .clone() - } - } - - struct NoopFlownode; - - #[async_trait] - impl Flownode for NoopFlownode { - async fn handle(&self, _request: FlowRequest) -> MetaResult { - unimplemented!() - } - - async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult { - unimplemented!() - } - - async fn handle_mark_window_dirty( - &self, - _req: DirtyWindowRequests, - ) -> MetaResult { - unimplemented!() - } - } - - #[async_trait] - impl FlownodeManager for ConcurrentMockNodeManager { - async fn flownode(&self, _node: &Peer) -> FlownodeRef { - Arc::new(NoopFlownode) - } - } - - struct RecordingFlownode { - requests_tx: mpsc::UnboundedSender, - } - - #[async_trait] - impl Flownode for RecordingFlownode { - async fn handle(&self, _request: FlowRequest) -> MetaResult { - unimplemented!() - } - - async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult { - unimplemented!() - } - - async fn handle_mark_window_dirty( - &self, - req: DirtyWindowRequests, - ) -> MetaResult { - self.requests_tx.send(req).unwrap(); - Ok(FlowResponse::default()) - } - } - - struct FlowNotificationMockNodeManager { - flownode: FlownodeRef, - } - - #[async_trait] - impl DatanodeManager for FlowNotificationMockNodeManager { - async fn datanode(&self, _node: &Peer) -> DatanodeRef { - unimplemented!() - } - } - - #[async_trait] - impl FlownodeManager for FlowNotificationMockNodeManager { - async fn flownode(&self, _node: &Peer) -> FlownodeRef { - self.flownode.clone() - } - } - - async fn mock_table_flownode_cache(table_id: TableId, peer: Peer) -> TableFlownodeSetCacheRef { - let cache = Arc::new(new_table_flownode_set_cache( - "test".to_string(), - CacheBuilder::new(1).build(), - Arc::new(MemoryKvBackend::default()), - )); - cache - .invalidate(&[CacheIdent::CreateFlow(CreateFlow { - flow_id: 1, - source_table_ids: vec![table_id], - partition_to_peer_mapping: vec![(0, peer.clone()), (1, peer)], - })]) - .await - .unwrap(); - cache - } - fn mock_flow_notification_sender( cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, - ) -> Notifier { + ) -> Notifier { let (tx, rx) = Notifier::try_new(16).unwrap(); start_flow_notification_worker(rx, cache, node_manager); tx @@ -2919,46 +2182,6 @@ mod tests { ); } - #[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) - } - } - - #[test] - fn test_batch_key_groups_by_skip_wal() { - let wal_ctx = session::context::QueryContext::arc(); - let skip_wal_ctx = session::context::QueryContext::arc(); - skip_wal_ctx.set_skip_wal(true); - let another_skip_wal_ctx = session::context::QueryContext::arc(); - another_skip_wal_ctx.set_skip_wal(true); - - let mut batches = HashMap::new(); - *batches.entry(batch_key_from_ctx(&wal_ctx)).or_insert(0) += 1; - *batches - .entry(batch_key_from_ctx(&skip_wal_ctx)) - .or_insert(0) += 1; - *batches - .entry(batch_key_from_ctx(&another_skip_wal_ctx)) - .or_insert(0) += 1; - *batches - .entry(batch_key_from_ctx(&session::context::QueryContext::arc())) - .or_insert(0) += 1; - - assert_eq!(batches.len(), 2); - assert_eq!(batches[&batch_key_from_ctx(&wal_ctx)], 2); - assert_eq!(batches[&batch_key_from_ctx(&skip_wal_ctx)], 2); - } - #[tokio::test] async fn test_cancelled_waiter_retains_request_slot_until_notification() { let limiter = RequestLimiter::try_new(1).unwrap(); @@ -3208,440 +2431,6 @@ mod tests { ); } - #[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_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![ - 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), - Field::new("region", ArrowDataType::Utf8, true), - ])); - let name_to_ids = - HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]); - let partition_columns = HashSet::from(["host"]); - - let (tag_columns, non_tag_indices) = - columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap(); - - assert_eq!(2, tag_columns.len()); - assert_eq!(&[0, 1, 2], non_tag_indices.as_slice()); - } - - #[test] - fn test_collect_tag_columns_and_non_tag_indices_prioritizes_essential_columns() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new("host", ArrowDataType::Utf8, true), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("region", ArrowDataType::Utf8, true), - ])); - let name_to_ids = - HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]); - let partition_columns = HashSet::from(["host", "region"]); - - let (_tag_columns, non_tag_indices): (_, SmallVec<[usize; 3]>) = - columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap(); - - assert_eq!(&[2, 1, 0, 3], non_tag_indices.as_slice()); - } - - #[test] - fn test_collect_tag_columns_and_non_tag_indices_rejects_unexpected_data_type() { - let schema = Arc::new(ArrowSchema::new(vec![ - 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), - Field::new("invalid", ArrowDataType::Boolean, true), - ])); - let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); - let partition_columns = HashSet::from(["host"]); - - let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); - - assert!(matches!( - result, - Err(Error::InvalidPromRemoteRequest { .. }) - )); - } - - #[test] - fn test_collect_tag_columns_and_non_tag_indices_rejects_int64_timestamp_column() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new("greptime_timestamp", ArrowDataType::Int64, false), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("host", ArrowDataType::Utf8, true), - ])); - let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); - let partition_columns = HashSet::from(["host"]); - - let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); - - assert!(matches!( - result, - Err(Error::InvalidPromRemoteRequest { .. }) - )); - } - - #[test] - fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_timestamp_column() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "ts1", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new( - "ts2", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("host", ArrowDataType::Utf8, true), - ])); - let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); - let partition_columns = HashSet::from(["host"]); - - let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); - - assert!(matches!( - result, - Err(Error::InvalidPromRemoteRequest { .. }) - )); - } - - #[test] - fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_value_column() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("value1", ArrowDataType::Float64, true), - Field::new("value2", ArrowDataType::Float64, true), - Field::new("host", ArrowDataType::Utf8, true), - ])); - let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); - let partition_columns = HashSet::from(["host"]); - - let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); - - assert!(matches!( - result, - Err(Error::InvalidPromRemoteRequest { .. }) - )); - } - - #[test] - fn test_modify_batch_sparse_with_taxonomy_per_batch() { - use arrow::array::BinaryArray; - use metric_engine::batch_modifier::modify_batch_sparse; - - let schema1 = Arc::new(ArrowSchema::new(vec![ - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("tag1", ArrowDataType::Utf8, true), - ])); - - let schema2 = Arc::new(ArrowSchema::new(vec![ - Field::new( - "greptime_timestamp", - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - Field::new("tag1", ArrowDataType::Utf8, true), - Field::new("tag2", ArrowDataType::Utf8, true), - ])); - let batch2 = RecordBatch::try_new( - schema2.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![2000])), - Arc::new(arrow::array::Float64Array::from(vec![2.0])), - Arc::new(StringArray::from(vec!["v1"])), - Arc::new(StringArray::from(vec!["v2"])), - ], - ) - .unwrap(); - - let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]); - let partition_columns = HashSet::new(); - - // A batch that only has tag1, same values as batch2 for ts and val. - let batch3 = RecordBatch::try_new( - schema1.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![2000])), - Arc::new(arrow::array::Float64Array::from(vec![2.0])), - Arc::new(StringArray::from(vec!["v1"])), - ], - ) - .unwrap(); - - // Simulate the new loop logic in flush_batch_physical: - // Resolve taxonomy FOR EACH BATCH. - let (tag_columns2, indices2) = - columns_taxonomy(&batch2.schema(), "table", &name_to_ids, &partition_columns).unwrap(); - let modified2 = modify_batch_sparse(batch2, 123, &tag_columns2, &indices2).unwrap(); - - let (tag_columns3, indices3) = - columns_taxonomy(&batch3.schema(), "table", &name_to_ids, &partition_columns).unwrap(); - let modified3 = modify_batch_sparse(batch3, 123, &tag_columns3, &indices3).unwrap(); - - let pk2 = modified2 - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let pk3 = modified3 - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - - // Now they SHOULD be different because tag2 is included in pk2 but not in pk3. - assert_ne!( - pk2.value(0), - pk3.value(0), - "PK should be different because batch2 has tag2!" - ); - } - - #[test] - fn test_transform_logical_batches_to_physical_success() { - let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); - - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 1, - batches: vec![batch], - row_count: 1, - }]; - - let name_to_ids = HashMap::from([("tag1".to_string(), 1)]); - let partition_columns = HashSet::new(); - let modified = - transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) - .unwrap(); - - assert_eq!(1, modified.len()); - assert_eq!(3, modified[0].num_columns()); - assert_eq!("__primary_key", modified[0].schema().field(0).name()); - assert_eq!("greptime_timestamp", modified[0].schema().field(1).name()); - assert_eq!("greptime_value", modified[0].schema().field(2).name()); - } - - #[test] - fn test_transform_logical_batches_to_physical_taxonomy_failure() { - let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); - - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 1, - batches: vec![batch], - row_count: 1, - }]; - - // tag1 is missing from name_to_ids, causing columns_taxonomy to fail. - let name_to_ids = HashMap::new(); - let partition_columns = HashSet::new(); - let err = - transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) - .unwrap_err(); - - assert!( - err.to_string() - .contains("not found in physical table column IDs") - ); - } - - #[test] - fn test_transform_logical_batches_to_physical_multiple_batches() { - let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); - let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0); - - let table_batches = vec![ - TableBatch { - table_name: "t1".to_string(), - table_id: 1, - batches: vec![batch1], - row_count: 1, - }, - TableBatch { - table_name: "t2".to_string(), - table_id: 2, - batches: vec![batch2], - row_count: 1, - }, - ]; - - let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]); - let partition_columns = HashSet::new(); - let modified = - transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) - .unwrap(); - - assert_eq!(2, modified.len()); - } - - #[test] - fn test_transform_logical_batches_to_physical_mixed_success_failure() { - let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); - let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0); - - let table_batches = vec![ - TableBatch { - table_name: "t1".to_string(), - table_id: 1, - batches: vec![batch1], - row_count: 1, - }, - TableBatch { - table_name: "t2".to_string(), - table_id: 2, - batches: vec![batch2], - row_count: 1, - }, - ]; - - // tag1 is missing from name_to_ids, causing batch1 to fail. - let name_to_ids = HashMap::from([("tag2".to_string(), 2)]); - let partition_columns = HashSet::new(); - let err = - transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) - .unwrap_err(); - - assert!(err.to_string().contains("tag1")); - } - #[tokio::test] async fn test_flush_batch_physical_uses_mockable_trait_dependencies() { let table_batches = vec![TableBatch { @@ -3801,88 +2590,31 @@ mod tests { } #[test] - fn test_plan_region_batches_splits_and_strips_partition_columns() { - let combined_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(), b"k2".as_slice()])), - Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])), - Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])), - Arc::new(StringArray::from(vec!["node-1", "node-2"])), - ], - ) - .unwrap(); - let mut planned_batches = plan_region_batches( - combined_batch, - 1024, - &TwoRegionPartitionRule { - partition_columns: vec!["host".to_string()], - }, - &["host".to_string()], - ) - .unwrap(); - planned_batches.sort_by_key(|planned| planned.region_id.region_number()); + fn test_batch_key_groups_by_skip_wal() { + use std::collections::HashMap; - 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!(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()); - } + use crate::batcher::logical_table::batch_key_from_ctx; - #[test] - fn test_encode_region_write_requests_builds_bulk_insert_requests() { - check_encode_region_write_requests(false); - check_encode_region_write_requests(true); - } + let wal_ctx = session::context::QueryContext::arc(); + let skip_wal_ctx = session::context::QueryContext::arc(); + skip_wal_ctx.set_skip_wal(true); + let another_skip_wal_ctx = session::context::QueryContext::arc(); + another_skip_wal_ctx.set_skip_wal(true); - fn check_encode_region_write_requests(skip_wal: bool) { - let planned_batch = PlannedRegionBatch { - region_id: RegionId::new(1024, 1), - 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), - ])), - vec![ - Arc::new(BinaryArray::from(vec![b"k1".as_slice()])), - Arc::new(TimestampMillisecondArray::from(vec![1000_i64])), - Arc::new(arrow::array::Float64Array::from(vec![1.0_f64])), - ], - ) - .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 mut batches = HashMap::new(); + *batches.entry(batch_key_from_ctx(&wal_ctx)).or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&skip_wal_ctx)) + .or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&another_skip_wal_ctx)) + .or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&session::context::QueryContext::arc())) + .or_insert(0) += 1; - assert_eq!(1, writes.len()); - assert_eq!(1, writes[0].datanode.id); - let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.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!(batches.len(), 2); + assert_eq!(batches[&batch_key_from_ctx(&wal_ctx)], 2); + assert_eq!(batches[&batch_key_from_ctx(&skip_wal_ctx)], 2); } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs new file mode 100644 index 00000000000..eba2335df03 --- /dev/null +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -0,0 +1,686 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; +use arrow::record_batch::RecordBatch; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; +use smallvec::SmallVec; +use snafu::{OptionExt, ResultExt, ensure}; +use table::metadata::TableId; + +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 { + /// The aligned logical record batch. + pub(in crate::batcher::logical_table) batch: RecordBatch, + /// The timestamp column index in `batch`. + pub(in crate::batcher::logical_table) timestamp_index: usize, +} + +impl RecordBatchWithTsIdx { + /// Creates a record batch with a validated timestamp column index. + pub fn try_new(batch: RecordBatch, timestamp_index: usize) -> Result { + let schema = batch.schema(); + let timestamp_field = schema.fields().get(timestamp_index).with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Timestamp column index {} is out of bounds for record batch with {} columns", + timestamp_index, + batch.num_columns() + ), + } + })?; + ensure!( + matches!(timestamp_field.data_type(), ArrowDataType::Timestamp(_, _)), + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Column at index {} is not a timestamp column: {:?}", + timestamp_index, + timestamp_field.data_type() + ), + } + ); + + Ok(Self { + batch, + timestamp_index, + }) + } + + #[cfg(test)] + pub(crate) fn into_parts(self) -> (RecordBatch, usize) { + (self.batch, self.timestamp_index) + } +} + +#[derive(Debug, Clone)] +pub struct TableBatch { + pub table_name: String, + pub table_id: TableId, + pub batches: Vec, + pub row_count: usize, +} + +/// Classifies columns in a logical-table batch for sparse primary-key conversion. +/// +/// Returns: +/// - `Vec`: all Utf8 tag columns sorted by tag name, used for +/// TSID and sparse primary-key encoding. +/// - `SmallVec<[usize; 3]>`: indices of columns copied into the physical batch +/// after `__primary_key`, ordered as `[greptime_timestamp, greptime_value, +/// partition_tag_columns...]`. +pub(in crate::batcher::logical_table) fn columns_taxonomy( + batch_schema: &Arc, + table_name: &str, + name_to_ids: &HashMap, + partition_columns: &HashSet<&str>, +) -> Result<(Vec, SmallVec<[usize; 3]>)> { + let mut tag_columns = Vec::new(); + let mut essential_column_indices = + SmallVec::<[usize; 3]>::with_capacity(2 + partition_columns.len()); + // Placeholder for greptime_timestamp and greptime_value + essential_column_indices.push(0); + essential_column_indices.push(0); + + let mut timestamp_index = None; + let mut value_index = None; + + for (index, field) in batch_schema.fields().iter().enumerate() { + match field.data_type() { + ArrowDataType::Utf8 => { + let column_id = name_to_ids.get(field.name()).copied().with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Column '{}' from logical table '{}' not found in physical table column IDs", + field.name(), + table_name + ), + } + })?; + tag_columns.push(TagColumnInfo { + name: field.name().clone(), + index, + column_id, + }); + + if partition_columns.contains(field.name().as_str()) { + essential_column_indices.push(index); + } + } + ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => { + ensure!( + timestamp_index.replace(index).is_none(), + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Duplicated timestamp column in logical table '{}' batch schema", + table_name + ), + } + ); + } + ArrowDataType::Float64 => { + ensure!( + value_index.replace(index).is_none(), + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Duplicated value column in logical table '{}' batch schema", + table_name + ), + } + ); + } + datatype => { + return error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Unexpected data type '{datatype:?}' in logical table '{}' batch schema", + table_name + ), + } + .fail(); + } + } + } + + let timestamp_index = + timestamp_index.with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Missing essential column '{}' in logical table '{}' batch schema", + greptime_timestamp(), + table_name + ), + })?; + let value_index = value_index.with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Missing essential column '{}' in logical table '{}' batch schema", + greptime_value(), + table_name + ), + })?; + + tag_columns.sort_by(|a, b| a.name.cmp(&b.name)); + + essential_column_indices[0] = timestamp_index; + essential_column_indices[1] = value_index; + + 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 +/// and applies sparse primary key modification. +pub(in crate::batcher::logical_table) fn transform_logical_batches_to_physical( + table_batches: &[TableBatch], + name_to_ids: &HashMap, + partition_columns_set: &HashSet<&str>, +) -> Result> { + let mut modified_batches: Vec = + Vec::with_capacity(table_batches.iter().map(|b| b.batches.len()).sum()); + + let mut modify_elapsed = Duration::ZERO; + let mut columns_taxonomy_elapsed = Duration::ZERO; + + for table_batch in table_batches { + let table_id = table_batch.table_id; + + for batch in &table_batch.batches { + let batch = &batch.batch; + let batch_schema = batch.schema(); + let start = Instant::now(); + let (tag_columns, essential_col_indices) = columns_taxonomy( + &batch_schema, + &table_batch.table_name, + name_to_ids, + partition_columns_set, + )?; + + columns_taxonomy_elapsed += start.elapsed(); + if tag_columns.is_empty() && essential_col_indices.is_empty() { + continue; + } + + let modified = { + let start = Instant::now(); + // The schema of modified batch is: __primary_key, timestamp, value, other partition columns... + let batch = modify_batch_sparse( + batch.clone(), + table_id, + &tag_columns, + &essential_col_indices, + )?; + modify_elapsed += start.elapsed(); + batch + }; + + modified_batches.push(modified); + } + } + + PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_physical_modify_batch"]) + .observe(modify_elapsed.as_secs_f64()); + PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_physical_columns_taxonomy"]) + .observe(columns_taxonomy_elapsed.as_secs_f64()); + + ensure!( + !modified_batches.is_empty(), + error::InternalSnafu { + err_msg: "No batches can be transformed during pending flush", + } + ); + Ok(modified_batches) +} + +/// Concatenates all modified batches into a single large batch. +/// +/// All modified batches share the same physical schema. +pub(in crate::batcher::logical_table) fn concat_modified_batches( + modified_batches: &[RecordBatch], +) -> Result { + let combined_schema = modified_batches[0].schema(); + let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_physical_concat_all"]) + .start_timer(); + concat_batches(&combined_schema, modified_batches).context(error::ArrowSnafu) +} + +#[cfg(test)] +mod tests { + use std::collections::{HashMap, HashSet}; + use std::sync::Arc; + + use arrow::array::{BinaryArray, StringArray, TimestampMillisecondArray}; + use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; + use arrow::record_batch::RecordBatch; + use smallvec::SmallVec; + + use crate::batcher::logical_table::batch_convert::{ + TableBatch, columns_taxonomy, strip_partition_columns_from_batch, + 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![ + 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), + Field::new("region", ArrowDataType::Utf8, true), + ])); + let name_to_ids = + HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]); + let partition_columns = HashSet::from(["host"]); + + let (tag_columns, non_tag_indices) = + columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap(); + + assert_eq!(2, tag_columns.len()); + assert_eq!(&[0, 1, 2], non_tag_indices.as_slice()); + } + + #[test] + fn test_collect_tag_columns_and_non_tag_indices_prioritizes_essential_columns() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("host", ArrowDataType::Utf8, true), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("region", ArrowDataType::Utf8, true), + ])); + let name_to_ids = + HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]); + let partition_columns = HashSet::from(["host", "region"]); + + let (_tag_columns, non_tag_indices): (_, SmallVec<[usize; 3]>) = + columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap(); + + assert_eq!(&[2, 1, 0, 3], non_tag_indices.as_slice()); + } + + #[test] + fn test_collect_tag_columns_and_non_tag_indices_rejects_unexpected_data_type() { + let schema = Arc::new(ArrowSchema::new(vec![ + 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), + Field::new("invalid", ArrowDataType::Boolean, true), + ])); + let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); + let partition_columns = HashSet::from(["host"]); + + let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); + + assert!(matches!( + result, + Err(Error::InvalidPromRemoteRequest { .. }) + )); + } + + #[test] + fn test_collect_tag_columns_and_non_tag_indices_rejects_int64_timestamp_column() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", ArrowDataType::Int64, false), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new("host", ArrowDataType::Utf8, true), + ])); + let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); + let partition_columns = HashSet::from(["host"]); + + let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); + + assert!(matches!( + result, + Err(Error::InvalidPromRemoteRequest { .. }) + )); + } + + #[test] + fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_timestamp_column() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts1", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new( + "ts2", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new("host", ArrowDataType::Utf8, true), + ])); + let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); + let partition_columns = HashSet::from(["host"]); + + let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); + + assert!(matches!( + result, + Err(Error::InvalidPromRemoteRequest { .. }) + )); + } + + #[test] + fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_value_column() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value1", ArrowDataType::Float64, true), + Field::new("value2", ArrowDataType::Float64, true), + Field::new("host", ArrowDataType::Utf8, true), + ])); + let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]); + let partition_columns = HashSet::from(["host"]); + + let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns); + + assert!(matches!( + result, + Err(Error::InvalidPromRemoteRequest { .. }) + )); + } + + #[test] + fn test_modify_batch_sparse_with_taxonomy_per_batch() { + use arrow::array::BinaryArray; + use metric_engine::batch_modifier::modify_batch_sparse; + + let schema1 = Arc::new(ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new("tag1", ArrowDataType::Utf8, true), + ])); + + let schema2 = Arc::new(ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new("tag1", ArrowDataType::Utf8, true), + Field::new("tag2", ArrowDataType::Utf8, true), + ])); + let batch2 = RecordBatch::try_new( + schema2.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![2000])), + Arc::new(arrow::array::Float64Array::from(vec![2.0])), + Arc::new(StringArray::from(vec!["v1"])), + Arc::new(StringArray::from(vec!["v2"])), + ], + ) + .unwrap(); + + let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]); + let partition_columns = HashSet::new(); + + // A batch that only has tag1, same values as batch2 for ts and val. + let batch3 = RecordBatch::try_new( + schema1.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![2000])), + Arc::new(arrow::array::Float64Array::from(vec![2.0])), + Arc::new(StringArray::from(vec!["v1"])), + ], + ) + .unwrap(); + + // Simulate the new loop logic in flush_batch_physical: + // Resolve taxonomy FOR EACH BATCH. + let (tag_columns2, indices2) = + columns_taxonomy(&batch2.schema(), "table", &name_to_ids, &partition_columns).unwrap(); + let modified2 = modify_batch_sparse(batch2, 123, &tag_columns2, &indices2).unwrap(); + + let (tag_columns3, indices3) = + columns_taxonomy(&batch3.schema(), "table", &name_to_ids, &partition_columns).unwrap(); + let modified3 = modify_batch_sparse(batch3, 123, &tag_columns3, &indices3).unwrap(); + + let pk2 = modified2 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let pk3 = modified3 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + // Now they SHOULD be different because tag2 is included in pk2 but not in pk3. + assert_ne!( + pk2.value(0), + pk3.value(0), + "PK should be different because batch2 has tag2!" + ); + } + + #[test] + fn test_transform_logical_batches_to_physical_success() { + let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); + + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 1, + batches: vec![batch], + row_count: 1, + }]; + + let name_to_ids = HashMap::from([("tag1".to_string(), 1)]); + let partition_columns = HashSet::new(); + let modified = + transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) + .unwrap(); + + assert_eq!(1, modified.len()); + assert_eq!(3, modified[0].num_columns()); + assert_eq!("__primary_key", modified[0].schema().field(0).name()); + assert_eq!("greptime_timestamp", modified[0].schema().field(1).name()); + assert_eq!("greptime_value", modified[0].schema().field(2).name()); + } + + #[test] + fn test_transform_logical_batches_to_physical_taxonomy_failure() { + let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); + + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 1, + batches: vec![batch], + row_count: 1, + }]; + + // tag1 is missing from name_to_ids, causing columns_taxonomy to fail. + let name_to_ids = HashMap::new(); + let partition_columns = HashSet::new(); + let err = + transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) + .unwrap_err(); + + assert!( + err.to_string() + .contains("not found in physical table column IDs") + ); + } + + #[test] + fn test_transform_logical_batches_to_physical_multiple_batches() { + let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); + let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0); + + let table_batches = vec![ + TableBatch { + table_name: "t1".to_string(), + table_id: 1, + batches: vec![batch1], + row_count: 1, + }, + TableBatch { + table_name: "t2".to_string(), + table_id: 2, + batches: vec![batch2], + row_count: 1, + }, + ]; + + let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]); + let partition_columns = HashSet::new(); + let modified = + transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) + .unwrap(); + + assert_eq!(2, modified.len()); + } + + #[test] + fn test_transform_logical_batches_to_physical_mixed_success_failure() { + let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0); + let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0); + + let table_batches = vec![ + TableBatch { + table_name: "t1".to_string(), + table_id: 1, + batches: vec![batch1], + row_count: 1, + }, + TableBatch { + table_name: "t2".to_string(), + table_id: 2, + batches: vec![batch2], + row_count: 1, + }, + ]; + + // tag1 is missing from name_to_ids, causing batch1 to fail. + let name_to_ids = HashMap::from([("tag2".to_string(), 2)]); + let partition_columns = HashSet::new(); + let err = + transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns) + .unwrap_err(); + + assert!(err.to_string().contains("tag1")); + } +} diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs new file mode 100644 index 00000000000..e3f817e70f2 --- /dev/null +++ b/src/servers/src/batcher/logical_table/region_write.rs @@ -0,0 +1,599 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use 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 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_telemetry::tracing_context::TracingContext; +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; + +#[derive(Clone)] +pub struct PhysicalTableMetadata { + pub table_info: TableInfoRef, + /// Mapping from column name to column id + pub col_name_to_ids: Option>, +} + +#[async_trait] +pub trait PhysicalFlushCatalogProvider: Send + Sync { + async fn physical_table( + &self, + catalog: &str, + schema: &str, + table_name: &str, + query_ctx: &session::context::QueryContext, + ) -> catalog::error::Result>; +} + +#[async_trait] +pub trait PhysicalFlushPartitionProvider: Send + Sync { + async fn find_table_partition_rule( + &self, + table_info: &TableInfo, + ) -> partition::error::Result; + + async fn find_region_leader(&self, region_id: RegionId) -> Result; +} + +#[async_trait] +pub trait PhysicalFlushNodeRequester: Send + Sync { + async fn handle( + &self, + peer: &Peer, + request: RegionRequest, + ) -> Result; +} + +#[derive(Clone)] +pub(in crate::batcher::logical_table) struct CatalogManagerPhysicalFlushAdapter { + pub(in crate::batcher::logical_table) catalog_manager: CatalogManagerRef, +} + +#[async_trait] +impl PhysicalFlushCatalogProvider for CatalogManagerPhysicalFlushAdapter { + async fn physical_table( + &self, + catalog: &str, + schema: &str, + table_name: &str, + query_ctx: &session::context::QueryContext, + ) -> catalog::error::Result> { + self.catalog_manager + .table(catalog, schema, table_name, Some(query_ctx)) + .await + .map(|table| { + table.map(|table| { + let table_info = table.table_info(); + let name_to_ids = table_info.name_to_ids(); + PhysicalTableMetadata { + table_info, + col_name_to_ids: name_to_ids, + } + }) + }) + } +} + +#[derive(Clone)] +pub(in crate::batcher::logical_table) struct PartitionManagerPhysicalFlushAdapter { + pub(in crate::batcher::logical_table) partition_manager: PartitionRuleManagerRef, +} + +#[async_trait] +impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter { + async fn find_table_partition_rule( + &self, + table_info: &TableInfo, + ) -> partition::error::Result { + 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) + } +} + +#[derive(Clone)] +pub(in crate::batcher::logical_table) struct NodeManagerPhysicalFlushAdapter { + pub(in crate::batcher::logical_table) node_manager: NodeManagerRef, +} + +#[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) + .await + .context(error::CommonMetaSnafu) + } +} + +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, +} + +#[cfg(test)] +impl PlannedRegionBatch { + pub(in crate::batcher::logical_table) 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 + .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, + physical_table_id: TableId, + region_number: u32, + mask: &partition::partition::RegionMask, +) -> Result> { + if mask.select_none() { + return Ok(None); + } + + let region_batch = if mask.select_all() { + stripped_batch.clone() + } else { + let _timer = 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)? + }; + + let row_count = region_batch.num_rows(); + if row_count == 0 { + return Ok(None); + } + + Ok(Some(PlannedRegionBatch { + region_id: RegionId::new(physical_table_id, region_number), + batch: region_batch, + })) +} + +pub(in crate::batcher::logical_table) 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)?; + + 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)? + { + planned_batches.push(planned_batch); + } + } + + Ok(planned_batches) +} + +pub(in crate::batcher::logical_table) async fn resolve_region_targets( + 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; + let (schema_bytes, data_header, payload) = { + let _timer = 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)? + }; + + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, + region_id: region_id.as_u64(), + partition_expr_version: None, + // 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, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: schema_bytes, + data_header, + payload, + })), + })), + }; + + region_writes.push(FlushRegionWrite { + datanode: resolved.datanode, + request, + }); + } + + Ok(region_writes) +} + +pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Error { + match error { + GrpcError::NotSupported { feat } => Error::NotSupported { feat }, + GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason }, + error => Error::Internal { + err_msg: error.to_string(), + }, + } +} + +#[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 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 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, + }; + use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager; + + struct TwoRegionPartitionRule { + partition_columns: Vec, + } + + impl PartitionRule for TwoRegionPartitionRule { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn partition_columns(&self) -> &[String] { + &self.partition_columns + } + + fn find_region( + &self, + _values: &[datatypes::prelude::Value], + ) -> partition::error::Result { + unimplemented!() + } + + fn split_record_batch( + &self, + _record_batch: &RecordBatch, + ) -> partition::error::Result> + { + Ok(HashMap::from([ + (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), + (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), + ( + 3, + RegionMask::new(BooleanArray::from(vec![false, false]), 0), + ), + ])) + } + } + + #[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!() + } + } + + #[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() { + let combined_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(), b"k2".as_slice()])), + Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])), + Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])), + Arc::new(StringArray::from(vec!["node-1", "node-2"])), + ], + ) + .unwrap(); + let mut planned_batches = plan_region_batches( + combined_batch, + 1024, + &TwoRegionPartitionRule { + partition_columns: vec!["host".to_string()], + }, + &["host".to_string()], + ) + .unwrap(); + planned_batches.sort_by_key(|planned| planned.region_id.region_number()); + + 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!(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()); + } + + #[test] + fn test_encode_region_write_requests_builds_bulk_insert_requests() { + check_encode_region_write_requests(false); + check_encode_region_write_requests(true); + } + + fn check_encode_region_write_requests(skip_wal: bool) { + let planned_batch = PlannedRegionBatch { + region_id: RegionId::new(1024, 1), + 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), + ])), + vec![ + Arc::new(BinaryArray::from(vec![b"k1".as_slice()])), + Arc::new(TimestampMillisecondArray::from(vec![1000_i64])), + Arc::new(arrow::array::Float64Array::from(vec![1.0_f64])), + ], + ) + .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(); + + assert_eq!(1, writes.len()); + assert_eq!(1, writes[0].datanode.id); + let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.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); + } +} diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs new file mode 100644 index 00000000000..b04883b27e0 --- /dev/null +++ b/src/servers/src/batcher/logical_table/test_util.rs @@ -0,0 +1,187 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use 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 arrow::array::{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::cache::{TableFlownodeSetCacheRef, new_table_flownode_set_cache}; +use common_meta::error::Result as MetaResult; +use common_meta::instruction::{CacheIdent, CreateFlow}; +use common_meta::kv_backend::memory::MemoryKvBackend; +use common_meta::node_manager::{ + DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, +}; +use moka::future::CacheBuilder; +use snafu::ResultExt; +use table::metadata::TableId; +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, + tag_value: &str, + ts: i64, + val: f64, +) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", ArrowDataType::Float64, true), + Field::new(tag_name, ArrowDataType::Utf8, true), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ts])), + Arc::new(arrow::array::Float64Array::from(vec![val])), + Arc::new(StringArray::from(vec![tag_value])), + ], + ) + .unwrap() +} + +pub(in crate::batcher::logical_table) fn mock_aligned_tag_batch( + tag_name: &str, + tag_value: &str, + ts: i64, + val: f64, +) -> RecordBatchWithTsIdx { + RecordBatchWithTsIdx::try_new(mock_tag_batch(tag_name, tag_value, ts, val), 0).unwrap() +} + +#[derive(Clone)] +pub(in crate::batcher::logical_table) struct ConcurrentMockNodeManager { + pub(in crate::batcher::logical_table) datanodes: Arc>, +} + +#[async_trait] +impl DatanodeManager for ConcurrentMockNodeManager { + async fn datanode(&self, node: &Peer) -> DatanodeRef { + self.datanodes + .get(&node.id) + .expect("datanode not found") + .clone() + } +} + +pub(in crate::batcher::logical_table) struct NoopFlownode; + +#[async_trait] +impl Flownode for NoopFlownode { + async fn handle(&self, _request: FlowRequest) -> MetaResult { + unimplemented!() + } + + async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult { + unimplemented!() + } + + async fn handle_mark_window_dirty( + &self, + _req: DirtyWindowRequests, + ) -> MetaResult { + unimplemented!() + } +} + +#[async_trait] +impl FlownodeManager for ConcurrentMockNodeManager { + async fn flownode(&self, _node: &Peer) -> FlownodeRef { + Arc::new(NoopFlownode) + } +} + +pub(in crate::batcher::logical_table) struct RecordingFlownode { + pub(in crate::batcher::logical_table) requests_tx: mpsc::UnboundedSender, +} + +#[async_trait] +impl Flownode for RecordingFlownode { + async fn handle(&self, _request: FlowRequest) -> MetaResult { + unimplemented!() + } + + async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult { + unimplemented!() + } + + async fn handle_mark_window_dirty(&self, req: DirtyWindowRequests) -> MetaResult { + self.requests_tx.send(req).unwrap(); + Ok(FlowResponse::default()) + } +} + +pub(in crate::batcher::logical_table) struct FlowNotificationMockNodeManager { + pub(in crate::batcher::logical_table) flownode: FlownodeRef, +} + +#[async_trait] +impl DatanodeManager for FlowNotificationMockNodeManager { + async fn datanode(&self, _node: &Peer) -> DatanodeRef { + unimplemented!() + } +} + +#[async_trait] +impl FlownodeManager for FlowNotificationMockNodeManager { + async fn flownode(&self, _node: &Peer) -> FlownodeRef { + self.flownode.clone() + } +} + +pub(in crate::batcher::logical_table) async fn mock_table_flownode_cache( + table_id: TableId, + peer: Peer, +) -> TableFlownodeSetCacheRef { + let cache = Arc::new(new_table_flownode_set_cache( + "test".to_string(), + CacheBuilder::new(1).build(), + Arc::new(MemoryKvBackend::default()), + )); + cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![table_id], + partition_to_peer_mapping: vec![(0, peer.clone()), (1, peer)], + })]) + .await + .unwrap(); + cache +} + +#[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) + } +} diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 27411a5f8a6..fd82079c016 100644 --- a/src/servers/src/http.rs +++ b/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>, + pending_rows_batcher: Option>, ) -> Self { let state = PromStoreState { prom_store_handler: handler, diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 60c000403fd..ed1c778f701 100644 --- a/src/servers/src/http/prom_store.rs +++ b/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>, + pub pending_rows_batcher: Option>, } #[derive(Debug, Serialize, Deserialize)] @@ -366,9 +366,9 @@ trait PromWriteBatcher: Send + Sync { } #[async_trait] -impl PromWriteBatcher for PendingRowsBatcher { +impl PromWriteBatcher for LogicalTablePendingRowsBatcher { async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { - PendingRowsBatcher::submit(self, requests, ctx).await + LogicalTablePendingRowsBatcher::submit(self, requests, ctx).await } } @@ -395,7 +395,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>, + pending_rows_batcher: Option>, prom_store_with_metric_engine: bool, mut batches: Vec, ) -> std::result::Result { @@ -461,7 +461,7 @@ async fn write_prometheus_rows_with_progress( async fn write_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, - pending_rows_batcher: Option>, + pending_rows_batcher: Option>, prom_store_with_metric_engine: bool, sample_batches: Vec, histogram_batches: Vec, diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index ac3ca064dfb..1106f89d3e9 100644 --- a/src/servers/src/lib.rs +++ b/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; diff --git a/src/servers/src/prom_row_builder.rs b/src/servers/src/prom_row_builder.rs index aad12814683..8b4e35bcd32 100644 --- a/src/servers/src/prom_row_builder.rs +++ b/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( diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index e0a749b813f..562bd48cdb0 100644 --- a/tests-integration/src/test_util.rs +++ b/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(), From 9bad25e7014165fd05496839f90ea975c5ebc9f9 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:43:54 +0000 Subject: [PATCH 30/63] ci: update dev-builder image tag (#9270) Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 9bed2ddc9b8..2526f3e31eb 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-7fa6f5f9-20260810091725 +DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-04f07a91-20260921031255 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu From 9e9a8cac2021cc2f20c58f5550647df76adcf18b Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 21 Sep 2026 03:58:38 +0000 Subject: [PATCH 31/63] fix(query): insert MergeScan into nested scalar subqueries (#9261) * fix(query): insert MergeScan into nested scalar subqueries DataFusion 55 keeps uncorrelated scalar subqueries as expression subqueries (enable_physical_uncorrelated_scalar_subquery, default true) instead of decorrelating them into joins, and executes them via the new physical ScalarSubqueryExec. DistPlannerAnalyzer::try_push_down walked the plan with a plain TreeNode transform that does not descend into expression subqueries, so MergeScan was only inserted for depth-1 subqueries. A scalar subquery nested inside another scalar subquery kept a bare frontend DistTable TableScan and failed at execution with "Unsupported operation: get stream from a distributed table". Use the subquery-aware transform so handle_subquery (PlanRewriter / MergeScan insertion) runs for subquery plans at every nesting depth. Fixes #9260. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): strengthen nested scalar subquery regression coverage Address review on #9261: - Replace the ineffective 'no bare TableScan' string check with a real subquery-aware plan walk (apply_with_subqueries); MergeScan hides its remote input from traversal, so any TableScan the walk reaches was genuinely left unwrapped. - Add a distributed regression case on a range-partitioned table so the nested inner aggregate must merge partial results across regions (global AVG feeding an outer SUM filter). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/analyzer.rs | 6 +- src/query/src/dist_plan/analyzer/test.rs | 74 +++++++++ .../distributed/explain/subqueries.result | 156 ++++++++++++++++++ .../cases/distributed/explain/subqueries.sql | 49 ++++++ .../standalone/common/subquery/table.result | 63 +++++++ .../standalone/common/subquery/table.sql | 23 +++ 6 files changed, 370 insertions(+), 1 deletion(-) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index d928671fc5d..2c1d066054a 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -247,7 +247,11 @@ fn unwrap_dictionary_literals(plan: LogicalPlan) -> DfResult DfResult { - let plan = plan.transform(&Self::inspect_plan_with_subquery)?; + // Use the subquery-aware transform so expression subqueries (including + // nested scalar subqueries left in place by DataFusion 55's + // `enable_physical_uncorrelated_scalar_subquery`) are visited at every + // depth and their table scans get wrapped in `MergeScan`. + let plan = plan.transform_with_subqueries(&Self::inspect_plan_with_subquery)?; let mut rewriter = PlanRewriter::new(&plan.data); let result = plan.data.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 36ea83fadae..554628b3de2 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -2982,3 +2982,77 @@ fn scheduled_none_falls_back_to_wall_clock() { "Remote should contain TimestampNanosecond:\n{result_str}" ); } + +/// Regression test for https://github.com/GreptimeTeam/greptimedb/issues/9260: +/// every scalar subquery level must get its own MergeScan, even when nested +/// inside another scalar subquery. +#[test] +fn expand_nested_scalar_subquery() { + init_default_ut_logging(); + let table_source = || { + Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(TestTable::table_with_name(0, "t".to_string())), + ))) as _ + }; + + // Innermost: SELECT max(t.number) FROM t + let innermost = LogicalPlanBuilder::scan_with_filters("t", table_source(), None, vec![]) + .unwrap() + .aggregate(Vec::::new(), vec![max(col("number"))]) + .unwrap() + .build() + .unwrap(); + // Middle: SELECT max(t.number) FROM t WHERE t.number > () + let middle = LogicalPlanBuilder::scan_with_filters("t", table_source(), None, vec![]) + .unwrap() + .filter(col("number").gt(Expr::ScalarSubquery(Subquery { + subquery: Arc::new(innermost), + outer_ref_columns: vec![], + spans: Default::default(), + }))) + .unwrap() + .aggregate(Vec::::new(), vec![max(col("number"))]) + .unwrap() + .build() + .unwrap(); + // Outermost projection: SELECT () AS m FROM t + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source(), None, vec![]) + .unwrap() + .project(vec![ + Expr::ScalarSubquery(Subquery { + subquery: Arc::new(middle), + outer_ref_columns: vec![], + spans: Default::default(), + }) + .alias("m"), + ]) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result_str = result.to_string(); + + // The main scan and both subquery levels must each be wrapped in a MergeScan. + let merge_scan_count = result_str + .matches("MergeScan [is_placeholder=false") + .count(); + assert_eq!( + 3, merge_scan_count, + "expected 3 MergeScan nodes (main + 2 subquery levels) in plan:\n{result_str}" + ); + // No bare TableScan may remain outside a MergeScan remote input. The + // subquery-aware walk still visits subquery plans, but cannot see into a + // MergeScan's hidden remote input (`MergeScanLogicalPlan::inputs()` is + // empty), so any TableScan it reaches was not wrapped. + result + .apply_with_subqueries(|node| { + assert!( + !matches!(node, LogicalPlan::TableScan(_)), + "unwrapped TableScan in plan:\n{result_str}" + ); + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); +} diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index 0ea758a6fe7..28fe75df516 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -652,3 +652,159 @@ DROP TABLE t2; Affected Rows: 0 +-- Regression test for https://github.com/GreptimeTeam/greptimedb/issues/9260: +-- nested scalar subqueries must get a MergeScan inside every subquery level. +CREATE TABLE nested_scalar (k INT, v BIGINT, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +INSERT INTO nested_scalar VALUES (1, 10, '2024-01-01 00:00:00'), (2, NULL, '2024-01-01 00:00:01'); + +Affected Rows: 2 + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: () AS m_| +|_|_Subquery:_| +|_|_Projection: ()_| +|_|_Subquery:_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Projection: max(nested_scalar.v)_| +|_|_Aggregate: groupBy=[[]], aggr=[[max(nested_scalar.v)]]_| +|_|_TableScan: nested_scalar_| +|_| ]]_| +|_|_EmptyRelation: rows=1_| +|_|_Projection:_| +|_|_Limit: skip=0, fetch=1_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| TableScan: nested_scalar_| +|_| ]]_| +| physical_plan | ScalarSubqueryExec: subqueries=1_| +|_|_ProjectionExec: expr=[scalar_subquery() as m]_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[]_| +|_|_GlobalLimitExec: skip=0, fetch=1_| +|_|_MergeScanExec: REDACTED +|_|_ScalarSubqueryExec: subqueries=1_| +|_|_ProjectionExec: expr=[scalar_subquery() as max(nested_scalar.v)] | +|_|_PlaceholderRowExec_| +|_|_CooperativeExec_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT (SELECT MAX(v) FROM nested_scalar WHERE v > (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + ++---+ +| m | ++---+ +| | ++---+ + +SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + ++----+ +| m | ++----+ +| 10 | ++----+ + +DROP TABLE nested_scalar; + +Affected Rows: 0 + +-- Same regression as above, but on a range-partitioned table so the inner +-- aggregate must merge partial results from multiple regions (the global +-- AVG must be computed across regions before the outer SUM filter runs). +CREATE TABLE nested_scalar_part (k INT, v BIGINT, ts TIMESTAMP TIME INDEX) +PARTITION ON COLUMNS (k) ( + k < 20, + k >= 20 AND k < 40, + k >= 40 +); + +Affected Rows: 0 + +INSERT INTO nested_scalar_part VALUES + (1, 10, '2024-01-01 00:00:00'), + (2, NULL, '2024-01-01 00:00:01'), + (21, 20, '2024-01-01 00:00:02'), + (25, 30, '2024-01-01 00:00:03'), + (41, 40, '2024-01-01 00:00:04'), + (50, 50, '2024-01-01 00:00:05'); + +Affected Rows: 6 + +-- global AVG(v) = 30, so the outer SUM(v) must be 40 + 50 = 90 +SELECT (SELECT SUM(v) FROM nested_scalar_part WHERE v > (SELECT AVG(v) FROM nested_scalar_part)) AS m FROM nested_scalar_part LIMIT 1; + ++----+ +| m | ++----+ +| 90 | ++----+ + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT (SELECT SUM(v) FROM nested_scalar_part WHERE v > (SELECT AVG(v) FROM nested_scalar_part)) AS m FROM nested_scalar_part LIMIT 1; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: () AS m_| +|_|_Subquery:_| +|_|_Aggregate: groupBy=[[]], aggr=[[sum(nested_scalar_part.v)]]_| +|_|_Filter: CAST(nested_scalar_part.v AS Float64) > ()_| +|_|_Subquery:_| +|_|_Aggregate: groupBy=[[]], aggr=[[__avg_merge(__avg_state(nested_scalar_part.v)) AS avg(nested_scalar_part.v)]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Aggregate: groupBy=[[]], aggr=[[__avg_state(CAST(nested_scalar_part.v AS Float64))]]_| +|_|_TableScan: nested_scalar_part_| +|_| ]]_| +|_|_Projection: nested_scalar_part.v_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| TableScan: nested_scalar_part_| +|_| ]]_| +|_|_Projection:_| +|_|_Limit: skip=0, fetch=1_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| TableScan: nested_scalar_part_| +|_| ]]_| +| physical_plan | ScalarSubqueryExec: subqueries=1_| +|_|_ProjectionExec: expr=[scalar_subquery() as m]_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_CoalescePartitionsExec: fetch=1_| +|_|_ProjectionExec: expr=[]_| +|_|_LocalLimitExec: fetch=1_| +|_|_MergeScanExec: REDACTED +|_|_ScalarSubqueryExec: subqueries=1_| +|_|_AggregateExec: mode=Final, gby=[], aggr=[sum(nested_scalar_part.v)]_| +|_|_CoalescePartitionsExec_| +|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(nested_scalar_part.v)]_| +|_|_FilterExec: CAST(v@0 AS Float64) > scalar_subquery()_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[v@1 as v]_| +|_|_MergeScanExec: REDACTED +|_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_merge(__avg_state(nested_scalar_part.v)) as avg(nested_scalar_part.v)]_| +|_|_CoalescePartitionsExec_| +|_|_AggregateExec: mode=Partial, gby=[], aggr=[__avg_merge(__avg_state(nested_scalar_part.v)) as avg(nested_scalar_part.v)] | +|_|_RepartitionExec: partitioning=REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +DROP TABLE nested_scalar_part; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/explain/subqueries.sql b/tests/cases/distributed/explain/subqueries.sql index 1935546f25d..ee5e0355f6e 100644 --- a/tests/cases/distributed/explain/subqueries.sql +++ b/tests/cases/distributed/explain/subqueries.sql @@ -202,3 +202,52 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x LIMIT 2; DROP TABLE t; DROP TABLE t1; DROP TABLE t2; + +-- Regression test for https://github.com/GreptimeTeam/greptimedb/issues/9260: +-- nested scalar subqueries must get a MergeScan inside every subquery level. +CREATE TABLE nested_scalar (k INT, v BIGINT, ts TIMESTAMP TIME INDEX); + +INSERT INTO nested_scalar VALUES (1, 10, '2024-01-01 00:00:00'), (2, NULL, '2024-01-01 00:00:01'); + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + +SELECT (SELECT MAX(v) FROM nested_scalar WHERE v > (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + +SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + +DROP TABLE nested_scalar; + +-- Same regression as above, but on a range-partitioned table so the inner +-- aggregate must merge partial results from multiple regions (the global +-- AVG must be computed across regions before the outer SUM filter runs). +CREATE TABLE nested_scalar_part (k INT, v BIGINT, ts TIMESTAMP TIME INDEX) +PARTITION ON COLUMNS (k) ( + k < 20, + k >= 20 AND k < 40, + k >= 40 +); + +INSERT INTO nested_scalar_part VALUES + (1, 10, '2024-01-01 00:00:00'), + (2, NULL, '2024-01-01 00:00:01'), + (21, 20, '2024-01-01 00:00:02'), + (25, 30, '2024-01-01 00:00:03'), + (41, 40, '2024-01-01 00:00:04'), + (50, 50, '2024-01-01 00:00:05'); + +-- global AVG(v) = 30, so the outer SUM(v) must be 40 + 50 = 90 +SELECT (SELECT SUM(v) FROM nested_scalar_part WHERE v > (SELECT AVG(v) FROM nested_scalar_part)) AS m FROM nested_scalar_part LIMIT 1; + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT (SELECT SUM(v) FROM nested_scalar_part WHERE v > (SELECT AVG(v) FROM nested_scalar_part)) AS m FROM nested_scalar_part LIMIT 1; + +DROP TABLE nested_scalar_part; diff --git a/tests/cases/standalone/common/subquery/table.result b/tests/cases/standalone/common/subquery/table.result index bcd107a35a5..24d447b0d67 100644 --- a/tests/cases/standalone/common/subquery/table.result +++ b/tests/cases/standalone/common/subquery/table.result @@ -182,3 +182,66 @@ SELECT * FROM (SELECT * FROM (SELECT 42 a), (SELECT 43 b)) JOIN (SELECT 44 c) ON | 42 | 43 | 44 | 45 | +----+----+----+----+ +-- skipped, unsupported feature: unnamed_subquery, see also: +-- https://github.com/GreptimeTeam/greptimedb/issues/5012 +-- SELECT * FROM (SELECT unnamed_subquery.a FROM (SELECT 42 a)), (SELECT unnamed_subquery.b FROM (SELECT 43 b)); +-- SELECT unnamed_subquery.a, unnamed_subquery2.b FROM (SELECT 42 a), (SELECT 43 b); +-- Regression test for https://github.com/GreptimeTeam/greptimedb/issues/9260: +-- nested scalar subqueries in the projection must not hit +-- "Unsupported operation: get stream from a distributed table". +CREATE TABLE nested_scalar (k INT, v BIGINT, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +INSERT INTO nested_scalar VALUES (1, 10, '2024-01-01 00:00:00'), (2, NULL, '2024-01-01 00:00:01'); + +Affected Rows: 2 + +-- control group: single-level scalar subqueries +SELECT (SELECT MAX(v) FROM nested_scalar) AS m FROM nested_scalar LIMIT 1; + ++----+ +| m | ++----+ +| 10 | ++----+ + +SELECT (SELECT MAX(v) FROM nested_scalar WHERE k = 1) AS m FROM nested_scalar LIMIT 1; + ++----+ +| m | ++----+ +| 10 | ++----+ + +-- scalar subquery nested inside another scalar subquery's filter +SELECT (SELECT MAX(v) FROM nested_scalar WHERE v > (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + ++---+ +| m | ++---+ +| | ++---+ + +-- scalar subquery nested one level deeper +SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + ++----+ +| m | ++----+ +| 10 | ++----+ + +-- three levels deep +SELECT (SELECT (SELECT (SELECT MAX(v) FROM nested_scalar))) AS m FROM nested_scalar LIMIT 1; + ++----+ +| m | ++----+ +| 10 | ++----+ + +DROP TABLE nested_scalar; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/subquery/table.sql b/tests/cases/standalone/common/subquery/table.sql index d49f3af739a..aeddda4a8a9 100644 --- a/tests/cases/standalone/common/subquery/table.sql +++ b/tests/cases/standalone/common/subquery/table.sql @@ -62,3 +62,26 @@ SELECT * FROM (SELECT * FROM (SELECT 42 a), (SELECT 43 b)) JOIN (SELECT 44 c) ON -- https://github.com/GreptimeTeam/greptimedb/issues/5012 -- SELECT * FROM (SELECT unnamed_subquery.a FROM (SELECT 42 a)), (SELECT unnamed_subquery.b FROM (SELECT 43 b)); -- SELECT unnamed_subquery.a, unnamed_subquery2.b FROM (SELECT 42 a), (SELECT 43 b); + +-- Regression test for https://github.com/GreptimeTeam/greptimedb/issues/9260: +-- nested scalar subqueries in the projection must not hit +-- "Unsupported operation: get stream from a distributed table". +CREATE TABLE nested_scalar (k INT, v BIGINT, ts TIMESTAMP TIME INDEX); + +INSERT INTO nested_scalar VALUES (1, 10, '2024-01-01 00:00:00'), (2, NULL, '2024-01-01 00:00:01'); + +-- control group: single-level scalar subqueries +SELECT (SELECT MAX(v) FROM nested_scalar) AS m FROM nested_scalar LIMIT 1; + +SELECT (SELECT MAX(v) FROM nested_scalar WHERE k = 1) AS m FROM nested_scalar LIMIT 1; + +-- scalar subquery nested inside another scalar subquery's filter +SELECT (SELECT MAX(v) FROM nested_scalar WHERE v > (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + +-- scalar subquery nested one level deeper +SELECT (SELECT (SELECT MAX(v) FROM nested_scalar)) AS m FROM nested_scalar LIMIT 1; + +-- three levels deep +SELECT (SELECT (SELECT (SELECT MAX(v) FROM nested_scalar))) AS m FROM nested_scalar LIMIT 1; + +DROP TABLE nested_scalar; From 6647dbc43738b231daf8448a56c80bad12737035 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 21 Sep 2026 04:14:22 +0000 Subject: [PATCH 32/63] feat!: upgrade DataFusion fork to 55.1.0 (#9177) Rebase the GreptimeDB DataFusion fork from official 55.0.0 to 55.1.0. DataFusion 55.1.0 is a patch release on branch-55 containing eleven cherry-picked fixes (schema-adaptation struct filters, cast/projection metadata propagation, nested-nullability aggregation adaptation, UnnestExec batch_size, RightMark join ordering panic, empty-struct ScalarValue, and FFI codec fixes). All twenty GreptimeDB fork patches rebase onto it with no textual or semantic overlap; none of them is absorbed upstream, so all are retained. Fork pin moves to discord9/datafusion branch greptimedb-55.1.0, commit 2aa87d52cdce7006af492330064738f33ed294c1 (55.1.0 + 20 patches). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 183 +++++++++++++++++++++++++++-------------------------- Cargo.toml | 58 ++++++++--------- 2 files changed, 121 insertions(+), 120 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2abfa9b7c37..7ebd6bdb3db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2239,7 +2239,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -2250,7 +2250,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -3771,8 +3771,8 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3823,8 +3823,8 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3847,8 +3847,8 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3870,8 +3870,8 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3896,8 +3896,8 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "futures", "log", @@ -3906,8 +3906,8 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-compression", @@ -3942,8 +3942,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3966,8 +3966,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3989,8 +3989,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4012,8 +4012,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4044,13 +4044,13 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" [[package]] name = "datafusion-execution" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4074,8 +4074,8 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4098,8 +4098,8 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4109,8 +4109,8 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4140,8 +4140,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4160,8 +4160,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4171,8 +4171,8 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-ord 59.2.0", @@ -4195,8 +4195,8 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4210,8 +4210,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4226,8 +4226,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4235,8 +4235,8 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "datafusion-doc", "quote", @@ -4245,8 +4245,8 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "chrono", @@ -4293,10 +4293,11 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", + "arrow-schema 59.2.0", "datafusion-common", "datafusion-expr", "datafusion-expr-common", @@ -4315,8 +4316,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4329,8 +4330,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "chrono", @@ -4346,8 +4347,8 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4365,8 +4366,8 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "arrow-data 59.2.0", @@ -4402,8 +4403,8 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-catalog", @@ -4428,8 +4429,8 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4438,8 +4439,8 @@ dependencies = [ [[package]] name = "datafusion-proto-models" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -4448,8 +4449,8 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4463,8 +4464,8 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow-schema 59.2.0", "async-trait", @@ -4477,8 +4478,8 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "arrow 59.2.0", "bigdecimal 0.4.8", @@ -4496,8 +4497,8 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "55.0.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=b66dca3260c8e04f314c8036cdfbb08f1a564d95#b66dca3260c8e04f314c8036cdfbb08f1a564d95" +version = "55.1.0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" dependencies = [ "async-recursion", "async-trait", @@ -5164,7 +5165,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -5764,7 +5765,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -6672,7 +6673,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -6743,7 +6744,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.61.2", + "windows-core 0.57.0", ] [[package]] @@ -7165,7 +7166,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -7203,7 +7204,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -9454,7 +9455,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 3.3.0", + "proc-macro-crate 1.3.1", "proc-macro2", "quote", "syn 2.0.117", @@ -11485,7 +11486,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11505,8 +11506,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11554,7 +11555,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11567,7 +11568,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -12016,7 +12017,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12981,7 +12982,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12994,7 +12995,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -13085,7 +13086,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -13963,7 +13964,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -15014,7 +15015,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -17244,7 +17245,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 54fddcce397..27e829a8cd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -130,22 +130,22 @@ crc32fast = "1" criterion = "0.7" crossbeam-utils = "0.8" dashmap = "6.1" -datafusion = "=55.0.0" -datafusion-common = "=55.0.0" -datafusion-datasource = "=55.0.0" -datafusion-expr = "=55.0.0" -datafusion-expr-common = "=55.0.0" -datafusion-functions = "=55.0.0" -datafusion-functions-aggregate-common = "=55.0.0" -datafusion-functions-window-common = "=55.0.0" -datafusion-optimizer = "=55.0.0" +datafusion = "=55.1.0" +datafusion-common = "=55.1.0" +datafusion-datasource = "=55.1.0" +datafusion-expr = "=55.1.0" +datafusion-expr-common = "=55.1.0" +datafusion-functions = "=55.1.0" +datafusion-functions-aggregate-common = "=55.1.0" +datafusion-functions-window-common = "=55.1.0" +datafusion-optimizer = "=55.1.0" datafusion-orc = "0.10.0" datafusion-pg-catalog = "0.18.3" -datafusion-physical-expr = "=55.0.0" -datafusion-physical-plan = "=55.0.0" -datafusion-proto = "=55.0.0" -datafusion-sql = "=55.0.0" -datafusion-substrait = "=55.0.0" +datafusion-physical-expr = "=55.1.0" +datafusion-physical-plan = "=55.1.0" +datafusion-proto = "=55.1.0" +datafusion-sql = "=55.1.0" +datafusion-substrait = "=55.1.0" datafusion_object_store = { package = "object_store", version = "0.13.2" } deadpool = "0.12" deadpool-postgres = "0.14" @@ -361,21 +361,21 @@ rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } From d375851abd3ed57035072aee62b1b7dae4ee4904 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 21 Sep 2026 05:33:42 +0000 Subject: [PATCH 33/63] fix(ci): stabilize long-range benchmark execution and artifact collection (#9241) * fix(ci): store long-range benchmark data on disk Signed-off-by: WenyXu * fix(ci): update vmbench runtime and stage benchmark artifacts Signed-off-by: WenyXu * fix(ci): reduce vmbench generator batches and parallelism Signed-off-by: WenyXu * fix(ci): use runtime with unified benchmark report metadata Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/workflows/vmbench-long-range.yml | 44 ++++++++++++++++++------ 1 file changed, 33 insertions(+), 11 deletions(-) diff --git a/.github/workflows/vmbench-long-range.yml b/.github/workflows/vmbench-long-range.yml index 484e2163a05..9d0088217ad 100644 --- a/.github/workflows/vmbench-long-range.yml +++ b/.github/workflows/vmbench-long-range.yml @@ -33,13 +33,18 @@ on: runtime_image: description: Aliyun runtime image (tag or digest reference) type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260917-65906ca9@sha256:5af45e298019474d4c5089d63031a15f9bceb663890dd87fa0a56876c99c8597 + default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260918-545dc570@sha256:0ff03e7e690575e875bb85b3054a03486dd9359073546dc715da177153cf7a29 required: true ecs_instance_type: description: ECS instance type (independent of Query Regression) type: string default: ecs.c9i.2xlarge required: true + benchmark_root: + description: Absolute disk-backed work directory (outside the o11ybench checkout) + type: string + default: /home/runner/benchmark-data + required: true system_disk_gib: description: System disk GiB (20..2048) type: number @@ -68,6 +73,7 @@ on: permissions: contents: read env: + BENCHMARK_ROOT: ${{ inputs.benchmark_root }} RANGE_PRESET: ${{ inputs.range_preset }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} VICTORIAMETRICS_TAG: ${{ inputs.victoriametrics_tag }} @@ -103,6 +109,10 @@ jobs: run: | set -euo pipefail [[ -n "$ECS_INSTANCE_TYPE" ]] + [[ "$BENCHMARK_ROOT" == /* && "$BENCHMARK_ROOT" != / && "$BENCHMARK_ROOT" != */ ]] + [[ "$BENCHMARK_ROOT" != /tmp && "$BENCHMARK_ROOT" != /tmp/* ]] + [[ "$BENCHMARK_ROOT" != *$'\n'* && "$BENCHMARK_ROOT" != *$'\r'* ]] + [[ "/${BENCHMARK_ROOT#/}/" != */../* && "/${BENCHMARK_ROOT#/}/" != */./* ]] for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers'; exit 1; } done @@ -136,7 +146,7 @@ jobs: # Keep the original warm defaults: serial 5/1 and c1/c4 30s/5s. docker run --rm "$resolved" scripts/o11ybench_plan.py "${args[@]}" \ --run-id "long-range-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT-warm" \ - --output-root "/tmp/o11ybench-long-range-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT-warm" \ + --output-root "$BENCHMARK_ROOT/o11ybench/o11ybench-long-range-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT-warm" \ --cache-policy warm > "$RUNNER_TEMP/vmbench-plans/warm.json" - uses: actions/upload-artifact@v4 with: @@ -185,10 +195,12 @@ jobs: RESOLVED_RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} RESOLVED_O11YBENCH_REF: ${{ needs.validate.outputs.o11ybench_ref }} RUN_ID: ${{ needs.validate.outputs.run_id }} - ARTIFACT_ROOT: /tmp/o11ybench-${{ needs.validate.outputs.run_id }} + ARTIFACT_ROOT: ${{ inputs.benchmark_root }}/o11ybench/o11ybench-${{ needs.validate.outputs.run_id }} METADATA_ROOT: ${{ github.workspace }}/vmbench-metadata O11YBENCH_DIR: ${{ github.workspace }}/o11ybench - TSDG_PARALLELISM: '8' + TSDG_PARALLELISM: '4' + TSDG_BATCH_SIZE: '100000' + O11YBENCH_GENERATOR_CPUS: '4' steps: - name: Checkout matching benchmark tools uses: actions/checkout@v4 @@ -209,9 +221,11 @@ jobs: test "$(git -C "$O11YBENCH_DIR" rev-parse HEAD)" = "$RESOLVED_O11YBENCH_REF" jq --version docker info - df -h /tmp /var/lib/docker + mkdir -p "$BENCHMARK_ROOT/o11ybench" + test "$(stat -f -c %T "$BENCHMARK_ROOT")" != tmpfs || { echo 'Benchmark root must be disk-backed'; exit 1; } + df -h "$BENCHMARK_ROOT" /var/lib/docker # Fail before generation if either dataset space or Docker storage is too small. - for path in /tmp /var/lib/docker; do + for path in "$BENCHMARK_ROOT" /var/lib/docker; do available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') ((available >= 10 * 1024 * 1024)) || { echo "Need at least 10 GiB free on $path"; exit 1; } done @@ -245,16 +259,24 @@ jobs: if: always() run: | docker ps -aq --filter "label=o11ybench.run-id=$RUN_ID-warm" | xargs -r docker rm -fv + - name: Stage benchmark artifacts + if: always() + run: | + mkdir -p "$RUNNER_TEMP/vmbench-artifacts" + if [[ -d "$METADATA_ROOT" ]]; then + rsync -a --no-owner --no-group "$METADATA_ROOT" "$RUNNER_TEMP/vmbench-artifacts/" + fi + if [[ -d "$ARTIFACT_ROOT-warm" ]]; then + # Exclude DB directories before traversal; they may be unreadable by the runner. + rsync -a --no-owner --no-group --exclude='db-data/' --exclude='/data/*.bin' \ + "$ARTIFACT_ROOT-warm/" "$RUNNER_TEMP/vmbench-artifacts/$(basename "$ARTIFACT_ROOT-warm")/" + fi - name: Upload HTML and results (no dataset or DB files) if: always() uses: actions/upload-artifact@v4 with: name: vmbench-long-range-${{ github.run_id }}-${{ github.run_attempt }} - path: | - ${{ env.METADATA_ROOT }} - ${{ env.ARTIFACT_ROOT }}-warm - !${{ env.ARTIFACT_ROOT }}-warm/data/*.bin - !${{ env.ARTIFACT_ROOT }}-warm/**/db-data/** + path: ${{ runner.temp }}/vmbench-artifacts if-no-files-found: warn retention-days: 14 teardown: From 4df557bb9efd37457bea28bf8755f9cbeff7cd44 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 21 Sep 2026 06:45:04 +0000 Subject: [PATCH 34/63] test: exclude testing feature completely (#9072) * test: exclude testing feature completely * chore: fmt --- tests-integration/Cargo.toml | 22 ++++++++++++++-------- tests/runner/Cargo.toml | 5 ++++- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index d1b0774af9e..d097432b828 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,7 +29,7 @@ cache.workspace = true catalog.workspace = true chrono.workspace = true clap.workspace = true -client = { workspace = true, features = ["testing"] } +client.workspace = true cmd.workspace = true common-base.workspace = true common-catalog.workspace = true @@ -40,7 +40,7 @@ common-event-recorder.workspace = true common-frontend.workspace = true common-grpc.workspace = true common-memory-manager.workspace = true -common-meta = { workspace = true, features = ["testing"] } +common-meta.workspace = true common-options.workspace = true common-procedure.workspace = true common-query.workspace = true @@ -56,7 +56,7 @@ datatypes.workspace = true dotenv.workspace = true flate2 = "1.0" flow.workspace = true -frontend = { workspace = true, features = ["testing"] } +frontend.workspace = true futures.workspace = true futures-util.workspace = true http.workspace = true @@ -64,16 +64,16 @@ hyper-util = { workspace = true, features = ["tokio"] } log-query = { workspace = true } loki-proto.workspace = true meta-client.workspace = true -meta-srv = { workspace = true, features = ["mock"] } -mito2 = { workspace = true, features = ["test-shared-fs-region-migration"] } +meta-srv.workspace = true +mito2.workspace = true object-store.workspace = true -operator = { workspace = true, features = ["testing"] } +operator.workspace = true plugins.workspace = true prost.workspace = true query.workspace = true rand.workspace = true serde_json.workspace = true -servers = { workspace = true, features = ["testing"] } +servers.workspace = true session.workspace = true snafu.workspace = true sql.workspace = true @@ -97,21 +97,26 @@ zstd.workspace = true [dev-dependencies] catalog = { workspace = true, features = ["testing"] } +client = { workspace = true, features = ["testing"] } +common-meta = { workspace = true, features = ["testing"] } common-procedure-test.workspace = true common-wal = { workspace = true, features = ["testing"] } datafusion.workspace = true datafusion-expr.workspace = true either.workspace = true +frontend = { workspace = true, features = ["testing"] } hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -mito2 = { workspace = true, features = ["test"] } +meta-srv = { workspace = true, features = ["mock"] } +mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", "default-rustls-ring", ] } opentelemetry-proto.workspace = true +operator = { workspace = true, features = ["testing"] } otel-arrow-rust.workspace = true partition.workspace = true paste.workspace = true @@ -122,6 +127,7 @@ rskafka.workspace = true rstest.workspace = true rstest_reuse.workspace = true sea-query.workspace = true +servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } similar-asserts.workspace = true store-api.workspace = true diff --git a/tests/runner/Cargo.toml b/tests/runner/Cargo.toml index 8d07c2ad1b3..2941e85fa5a 100644 --- a/tests/runner/Cargo.toml +++ b/tests/runner/Cargo.toml @@ -10,7 +10,7 @@ workspace = true [dependencies] async-trait.workspace = true clap.workspace = true -client = { workspace = true, features = ["testing"] } +client.workspace = true common-error.workspace = true common-query.workspace = true common-recordbatch.workspace = true @@ -34,3 +34,6 @@ tokio.workspace = true tokio-postgres = { workspace = true } tokio-stream.workspace = true toml.workspace = true + +[dev-dependencies] +client = { workspace = true, features = ["testing"] } From 95f693c3966f18f9c2044d342896625ae0c59bb2 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 21 Sep 2026 07:08:39 +0000 Subject: [PATCH 35/63] perf(servers): stream Prometheus HTTP responses incrementally (#9229) * perf(servers): stream Prometheus HTTP responses incrementally Replace try_collect + full-buffer conversion in the Stream branch of from_query_result with incremental per-batch merging via the shared merge_batch path, so record batches are dropped as soon as they are consumed instead of all coexisting in memory. - Add ColumnLayout to dedupe schema inference shared by both paths - Add SeriesKeyLookup (Equivalent-based borrow lookup) so series keys are looked up by borrowed label slices without materializing owned keys per lookup - Empty streams short-circuit to empty data before schema inference - Add stream-side mixed histogram + Vector dual-path equivalence test Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(servers): address review comments on stream response tests - Inline single-caller tags_capacity helper - Drop timestamp column from empty-stream test so it pins the short-circuit happens before schema inference - Fix cross-batch merge test data so a series actually spans the batch boundary (second batch carries earlier timestamps) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(servers): add lazy stream release assertion and pinned cross-batch content - LazyDropCheckingStream creates one batch per poll and asserts via Weak that the previous batch was released before the next poll, protecting the core property that batches are dropped eagerly - Pin series "a" merged samples to exact timestamps and values in stream_result_matches_record_batches_result, so a regression in shared merge_batch logic cannot make both paths agree on wrong output Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/http/result/prometheus_resp.rs | 820 +++++++++++++----- 1 file changed, 612 insertions(+), 208 deletions(-) diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index eca38338719..518d0f65e32 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -36,12 +36,14 @@ use common_query::promql_annotations::{ PromqlAnnotationCollector, get_promql_annotation_collector, }; use common_query::{Output, OutputData}; -use common_recordbatch::RecordBatches; +use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::ConcreteDataType; -use indexmap::IndexMap; +use datatypes::schema::Schema; +use futures::TryStreamExt; use indexmap::map::RawEntryApiV1; use indexmap::map::raw_entry_v1::RawEntryMut; +use indexmap::{Equivalent, IndexMap}; use itertools::Either; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; @@ -236,16 +238,9 @@ impl PrometheusJsonResponse { OutputData::RecordBatches(batches) => Self::success( Self::record_batches_to_data(batches, metric_name, result_type)?, ), - OutputData::Stream(stream) => { - let record_batches = RecordBatches::try_collect(stream) - .await - .context(CollectRecordbatchSnafu)?; - Self::success(Self::record_batches_to_data( - record_batches, - metric_name, - result_type, - )?) - } + OutputData::Stream(stream) => Self::success( + Self::consume_stream_to_data(stream, metric_name, result_type).await?, + ), OutputData::AffectedRows(_) => Self::error( StatusCode::Unexpected, "expected data result, but got affected rows", @@ -299,22 +294,83 @@ impl PrometheusJsonResponse { ) -> Result { // Return empty result if no batches if batches.iter().next().is_none() { - return Ok(PrometheusResponse::PromData(PromData { - result_type: result_type.to_string(), - ..Default::default() - })); + return Ok(empty_data(result_type)); } - // infer semantic type of each column from schema. - // TODO(ruihang): wish there is a better way to do this. + let layout = ColumnLayout::infer(&batches.schema())?; + + // Preserves the order of output tags. + // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. + let mut buffer = IndexMap::, PromSeriesSamples>::new(); + + for batch in batches.iter() { + merge_batch(&mut buffer, &layout, batch, metric_name.as_deref())?; + } + + samples_to_data(buffer, result_type) + } + + /// Consume a streaming query result into [PromData]. + /// + /// Every batch is merged into the per-series buffer and dropped before the + /// next one is polled, so the whole query result and the assembled series + /// never coexist in memory. + async fn consume_stream_to_data( + mut stream: SendableRecordBatchStream, + metric_name: Option, + result_type: ValueType, + ) -> Result { + // An empty stream carries no batch to infer the column roles from, and + // reads as an empty result, like empty [RecordBatches]. + let Some(first_batch) = stream.try_next().await.context(CollectRecordbatchSnafu)? else { + return Ok(empty_data(result_type)); + }; + + let layout = ColumnLayout::infer(&stream.schema())?; + + // Preserves the order of output tags. + // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. + let mut buffer = IndexMap::, PromSeriesSamples>::new(); + + merge_batch(&mut buffer, &layout, &first_batch, metric_name.as_deref())?; + // Release the batch before polling the next one. + drop(first_batch); + + while let Some(batch) = stream.try_next().await.context(CollectRecordbatchSnafu)? { + merge_batch(&mut buffer, &layout, &batch, metric_name.as_deref())?; + } + + samples_to_data(buffer, result_type) + } +} + +/// The empty data of a Prometheus HTTP API response of `result_type`. +fn empty_data(result_type: ValueType) -> PrometheusResponse { + PrometheusResponse::PromData(PromData { + result_type: result_type.to_string(), + ..Default::default() + }) +} + +/// The role of every column of a query result. +struct ColumnLayout { + timestamp_column_index: usize, + tag_column_indices: Vec, + tag_names: Vec, + first_field_column_index: Option, + native_histogram_column_index: Option, +} + +impl ColumnLayout { + /// Infer the role of every column from `schema`. + // TODO(ruihang): wish there is a better way to do this. + fn infer(schema: &Schema) -> Result { let mut timestamp_column_index = None; let mut tag_column_indices = Vec::new(); let mut first_field_column_index = None; let mut native_histogram_column_index = None; - let mut num_label_columns = 0; - - for (i, column) in batches.schema().column_schemas().iter().enumerate() { + for (i, column) in schema.column_schemas().iter().enumerate() { match column.data_type { ConcreteDataType::Timestamp(datatypes::types::TimestampType::Millisecond(_)) if timestamp_column_index.is_none() => @@ -343,7 +399,6 @@ impl PrometheusJsonResponse { } ConcreteDataType::String(_) => { tag_column_indices.push(i); - num_label_columns += 1; } _ => {} } @@ -359,215 +414,255 @@ impl PrometheusJsonResponse { .fail(); } - // Preserves the order of output tags. - // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. - let mut buffer = IndexMap::, PromSeriesSamples>::new(); + let tag_names = tag_column_indices + .iter() + .map(|index| schema.column_name_by_index(*index).to_string()) + .collect::>(); - // Only a series that is new to `buffer` needs its own key vector. - let mut tags = Vec::with_capacity(num_label_columns + 1); + Ok(Self { + timestamp_column_index, + tag_column_indices, + tag_names, + first_field_column_index, + native_histogram_column_index, + }) + } +} - let schema = batches.schema(); - for batch in batches.iter() { - // prepare things... - let tag_columns = tag_column_indices +/// The lookup key of a series in the per-series sample buffer. +/// +/// The blanket `Equivalent` implementation compares through `Borrow`, which +/// `Vec<(String, String)>` cannot provide for labels borrowed from a batch, so +/// this wrapper compares the labels itself. It is built for a single lookup and +/// hashes like the owned key it looks up. +struct SeriesKeyLookup<'a>(&'a [(&'a str, &'a str)]); + +impl Equivalent> for SeriesKeyLookup<'_> { + fn equivalent(&self, key: &Vec<(String, String)>) -> bool { + self.0.len() == key.len() + && self + .0 .iter() - .map(|i| batch.column(*i)) - .collect::>(); - let tag_names = tag_column_indices - .iter() - .map(|c| schema.column_name_by_index(*c)) - .collect::>(); - let timestamp_column = batch - .column(timestamp_column_index) - .as_primitive::(); + .zip(key) + .all(|((name, value), (key_name, key_value))| { + *name == key_name.as_str() && *value == key_value.as_str() + }) + } +} - let field_array = first_field_column_index - .map(|index| arrow::compute::cast(batch.column(index), &DataType::Float64)) - .transpose() - .context(ArrowSnafu)?; - let field_column = field_array - .as_ref() - .map(|array| array.as_primitive::()); - let native_histogram_column = native_histogram_column_index - .map(|index| { - batch - .column(index) - .as_any() - .downcast_ref::() - .with_context(|| UnexpectedResultSnafu { - reason: "native histogram column is not a struct array", - }) +/// Merge the rows of one batch into the per-series sample buffer. +fn merge_batch( + buffer: &mut IndexMap, PromSeriesSamples>, + layout: &ColumnLayout, + batch: &RecordBatch, + metric_name: Option<&str>, +) -> Result<()> { + // Scratch buffer of the labels of the series being read, borrowed from + // `batch` and `layout` and refilled for every series that is not in + // `buffer` yet, so looking a series up allocates nothing and only a series + // that is new allocates its labels. + let mut tags: Vec<(&str, &str)> = Vec::with_capacity(layout.tag_column_indices.len() + 1); + + // prepare things... + let tag_columns = layout + .tag_column_indices + .iter() + .map(|i| batch.column(*i)) + .collect::>(); + let timestamp_column = batch + .column(layout.timestamp_column_index) + .as_primitive::(); + + let field_array = layout + .first_field_column_index + .map(|index| arrow::compute::cast(batch.column(index), &DataType::Float64)) + .transpose() + .context(ArrowSnafu)?; + let field_column = field_array + .as_ref() + .map(|array| array.as_primitive::()); + let native_histogram_column = layout + .native_histogram_column_index + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .with_context(|| UnexpectedResultSnafu { + reason: "native histogram column is not a struct array", + }) + }) + .transpose()?; + + // Read the labels once per run of rows that share them instead of + // once per row. `label_runs` finds every boundary, so the probe only + // decides whether looking for runs is worth its cost. + let label_runs = if prefer_label_runs(&tag_columns, batch.num_rows()) { + Either::Left(label_runs(&tag_columns, batch.num_rows())?.into_iter()) + } else { + Either::Right((0..batch.num_rows()).map(|row| row..row + 1)) + }; + + // assemble rows + for run in label_runs { + let mut run_entry_index = None; + for row_index in run { + let value = field_column.and_then(|field_column| { + if !field_column.is_valid(row_index) { + return None; + } + let value = field_column.value(row_index); + (!is_prometheus_stale_nan(value)) + .then_some((timestamp_column.value(row_index), value)) + }); + let histogram = native_histogram_column + .and_then(|column| { + read_histogram(column, row_index) + .context(DataFusionSnafu) + .transpose() + }) + .transpose()? + .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) + .map(|histogram| { + prometheus_native_histogram(&histogram) + .map(|histogram| (timestamp_column.value(row_index), histogram)) }) .transpose()?; - // Read the labels once per run of rows that share them instead of - // once per row. `label_runs` finds every boundary, so the probe only - // decides whether looking for runs is worth its cost. - let label_runs = if prefer_label_runs(&tag_columns, batch.num_rows()) { - Either::Left(label_runs(&tag_columns, batch.num_rows())?.into_iter()) - } else { - Either::Right((0..batch.num_rows()).map(|row| row..row + 1)) - }; + if value.is_none() && histogram.is_none() { + continue; + } - // assemble rows - for run in label_runs { - let mut run_entry_index = None; - for row_index in run { - let value = field_column.and_then(|field_column| { - if !field_column.is_valid(row_index) { - return None; + let entry_index = match run_entry_index { + Some(index) => index, + None => { + // retrieve tags + tags.clear(); + if let Some(metric_name) = metric_name { + tags.push((METRIC_NAME, metric_name)); + } + for (tag_column, tag_name) in tag_columns.iter().zip(layout.tag_names.iter()) { + if let Some(tag_value) = string_array_value_at_index(tag_column, row_index) + { + tags.push((tag_name.as_str(), tag_value)); } - let value = field_column.value(row_index); - (!is_prometheus_stale_nan(value)) - .then_some((timestamp_column.value(row_index), value)) - }); - let histogram = native_histogram_column - .and_then(|column| { - read_histogram(column, row_index) - .context(DataFusionSnafu) - .transpose() - }) - .transpose()? - .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) - .map(|histogram| { - prometheus_native_histogram(&histogram) - .map(|histogram| (timestamp_column.value(row_index), histogram)) - }) - .transpose()?; - - if value.is_none() && histogram.is_none() { - continue; } - let entry_index = match run_entry_index { - Some(index) => index, - None => { - // retrieve tags - tags.clear(); - if let Some(metric_name) = &metric_name { - tags.push((METRIC_NAME, metric_name.as_str())); - } - for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { - if let Some(tag_value) = - string_array_value_at_index(tag_column, row_index) - { - tags.push((tag_name, tag_value)); - } - } - - let hash = buffer.hasher().hash_one(&tags); - let entry = buffer - .raw_entry_mut_v1() - .from_key_hashed_nocheck(hash, &tags); - let index = entry.index(); - if let RawEntryMut::Vacant(entry) = entry { - // Hand the key buffer over with the hash that was - // just computed for the lookup. - let key = std::mem::replace( - &mut tags, - Vec::with_capacity(num_label_columns + 1), - ); - entry.insert_hashed_nocheck( - hash, - key, - PromSeriesSamples::default(), - ); - } - run_entry_index = Some(index); - index - } - }; - let samples = &mut buffer[entry_index]; - if let Some((timestamp_millis, histogram)) = histogram { - samples - .histograms - .push((timestamp_millis as f64 / 1000.0, histogram)); - } else if let Some((timestamp_millis, value)) = value { - samples.values.push(( - timestamp_millis as f64 / 1000.0, - PromSampleValue::Number(value), - )); + // Borrowed and owned labels hash the same: both + // hash their `str` contents, so this hash serves the + // lookup and the insertion of a new series below. + let hash = buffer.hasher().hash_one(&tags); + let entry = buffer + .raw_entry_mut_v1() + .from_key_hashed_nocheck(hash, &SeriesKeyLookup(tags.as_slice())); + let index = entry.index(); + if let RawEntryMut::Vacant(entry) = entry { + // A series that is new to the buffer takes + // ownership of its labels. + let key = tags + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect::>(); + entry.insert_hashed_nocheck(hash, key, PromSeriesSamples::default()); } + run_entry_index = Some(index); + index } + }; + let samples = &mut buffer[entry_index]; + if let Some((timestamp_millis, histogram)) = histogram { + samples + .histograms + .push((timestamp_millis as f64 / 1000.0, histogram)); + } else if let Some((timestamp_millis, value)) = value { + samples.values.push(( + timestamp_millis as f64 / 1000.0, + PromSampleValue::Number(value), + )); } } + } + Ok(()) +} - // initialize result to return - let mut result = match result_type { - ValueType::Vector => PromQueryResult::Vector(vec![]), - ValueType::Matrix => PromQueryResult::Matrix(vec![]), - ValueType::Scalar => PromQueryResult::Scalar(None), - ValueType::String => PromQueryResult::String(None), - }; +/// Assemble the per-series samples of `buffer` into the data of a Prometheus +/// HTTP API response of `result_type`. +fn samples_to_data( + buffer: IndexMap, PromSeriesSamples>, + result_type: ValueType, +) -> Result { + // initialize result to return + let mut result = match result_type { + ValueType::Vector => PromQueryResult::Vector(vec![]), + ValueType::Matrix => PromQueryResult::Matrix(vec![]), + ValueType::Scalar => PromQueryResult::Scalar(None), + ValueType::String => PromQueryResult::String(None), + }; - // accumulate data into result - buffer.into_iter().for_each(|(tags, mut samples)| { - let metric = tags - .into_iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect::>(); - match result { - PromQueryResult::Vector(ref mut v) => { - let histogram = samples.histograms.pop(); - let value = if histogram.is_none() { - samples - .values - .pop() - .map(|(timestamp, value)| (timestamp, value.into_string())) - } else { - None - }; - v.push(PromSeriesVector { - metric, - value, - histogram, - }); - } - PromQueryResult::Matrix(ref mut v) => { - // sort values by timestamp - if !samples.values.is_sorted_by(|a, b| a.0 <= b.0) { - samples - .values - .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); - } - if !samples.histograms.is_sorted_by(|a, b| a.0 <= b.0) { - samples - .histograms - .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); - } - - v.push(PromSeriesMatrix { - metric, - values: samples.values, - histograms: samples.histograms, - }); - } - PromQueryResult::Scalar(ref mut v) => { - *v = samples + // accumulate data into result + buffer.into_iter().for_each(|(tags, mut samples)| { + let metric = tags.into_iter().collect::>(); + match result { + PromQueryResult::Vector(ref mut v) => { + let histogram = samples.histograms.pop(); + let value = if histogram.is_none() { + samples .values .pop() - .map(|(timestamp, value)| (timestamp, value.into_string())); - } - PromQueryResult::String(ref mut _v) => { - // TODO(ruihang): Not supported yet - } + .map(|(timestamp, value)| (timestamp, value.into_string())) + } else { + None + }; + v.push(PromSeriesVector { + metric, + value, + histogram, + }); } - }); + PromQueryResult::Matrix(ref mut v) => { + // sort values by timestamp + if !samples.values.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .values + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + } + if !samples.histograms.is_sorted_by(|a, b| a.0 <= b.0) { + samples + .histograms + .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + } - // sort matrix by metric - // see: https://prometheus.io/docs/prometheus/3.5/querying/api/#range-vectors - if let PromQueryResult::Matrix(ref mut v) = result { - v.sort_by(|a, b| a.metric.cmp(&b.metric)); + v.push(PromSeriesMatrix { + metric, + values: samples.values, + histograms: samples.histograms, + }); + } + PromQueryResult::Scalar(ref mut v) => { + *v = samples + .values + .pop() + .map(|(timestamp, value)| (timestamp, value.into_string())); + } + PromQueryResult::String(ref mut _v) => { + // TODO(ruihang): Not supported yet + } } + }); - let result_type_string = result_type.to_string(); - let data = PrometheusResponse::PromData(PromData { - result_type: result_type_string, - result, - }); - - Ok(data) + // sort matrix by metric + // see: https://prometheus.io/docs/prometheus/3.5/querying/api/#range-vectors + if let PromQueryResult::Matrix(ref mut v) = result { + v.sort_by(|a, b| a.metric.cmp(&b.metric)); } + + let result_type_string = result_type.to_string(); + let data = PrometheusResponse::PromData(PromData { + result_type: result_type_string, + result, + }); + + Ok(data) } /// Ranges of consecutive rows whose label values are all equal. @@ -660,13 +755,322 @@ mod tests { }; use common_recordbatch::{RecordBatch, RecordBatches}; use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::vectors::{ Float64Vector, StringVector, StructVector, TimestampMillisecondVector, VectorRef, }; use super::*; + #[tokio::test] + async fn empty_stream_reads_as_empty_data() { + // No timestamp column on purpose: if the empty-stream short-circuit ran + // after `ColumnLayout::infer`, inference would fail with + // "no timestamp column found" instead of returning empty data. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let stream = RecordBatches::try_new(schema, vec![]).unwrap().as_stream(); + + let response = + PrometheusJsonResponse::consume_stream_to_data(stream, None, ValueType::Matrix).await; + + assert!(matches!( + response.unwrap(), + PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) if series.is_empty() + )); + } + + #[tokio::test] + async fn stream_result_matches_record_batches_result() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + // Two batches, with series "a" split across them: batch 0 holds its + // later samples, batch 1 its earlier ones plus series "b", so streaming + // must merge samples of one series across a batch boundary. + let batch_rows: [Vec<(&str, i64, f64)>; 2] = [ + vec![ + ("a", 4_000, 0.0), + ("a", 5_000, 1.0), + ("a", 6_000, 2.0), + ("a", 7_000, 3.0), + ], + vec![ + ("a", 0, 4.0), + ("a", 1_000, 5.0), + ("b", 2_000, 6.0), + ("b", 3_000, 7.0), + ], + ]; + let batches: Vec<_> = batch_rows + .into_iter() + .map(|rows| { + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values( + rows.iter().map(|(_, ts, _)| *ts), + )) as _, + Arc::new(StringVector::from( + rows.iter() + .map(|(host, _, _)| Some(*host)) + .collect::>(), + )) as _, + Arc::new(Float64Vector::from_values(rows.iter().map(|(_, _, v)| *v))) as _, + ], + ) + .unwrap() + }) + .collect(); + + let streamed = PrometheusJsonResponse::consume_stream_to_data( + RecordBatches::try_new(schema.clone(), batches.clone()) + .unwrap() + .as_stream(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .await + .unwrap(); + let eager = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .unwrap(); + + assert_eq!( + serde_json::to_value(&streamed).unwrap(), + serde_json::to_value(&eager).unwrap() + ); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = streamed + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), 2); + // Pin the merged content of the cross-batch series, not just the + // two-path agreement: both paths share merge_batch, so a regression + // there could make them agree on the same wrong output. + let series_a = series + .iter() + .find(|s| s.metric.get("host").map(String::as_str) == Some("a")) + .unwrap(); + let samples_a: Vec<(i64, f64)> = series_a + .values + .iter() + .map(|(ts, value)| { + let PromSampleValue::Number(v) = value else { + panic!("expected numeric sample"); + }; + ((ts * 1000.0) as i64, *v) + }) + .collect(); + assert_eq!( + samples_a, + vec![ + (0, 4.0), + (1_000, 5.0), + (4_000, 0.0), + (5_000, 1.0), + (6_000, 2.0), + (7_000, 3.0), + ] + ); + } + + #[tokio::test] + async fn stream_mixed_histogram_result_matches_record_batches_result() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("kind", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("float", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("histogram", native_histogram_value_type().clone(), true), + ])); + // The float series splits across the batch boundary so streaming must + // merge its samples like the eager path; the histogram row only arrives + // in the second batch. + let first = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values([1_000, 2_000])) as _, + Arc::new(StringVector::from(vec![Some("float"); 2])) as _, + Arc::new(Float64Vector::from(vec![Some(1.25), Some(0.5)])) as _, + histogram_vector(&[None, None]), + ], + ) + .unwrap(); + let second = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values([3_000, 3_000])) as _, + Arc::new(StringVector::from(vec![Some("float"), Some("histogram")])) as _, + Arc::new(Float64Vector::from(vec![Some(2.0), None])) as _, + histogram_vector(&[None, Some(sample_histogram())]), + ], + ) + .unwrap(); + + let streamed = PrometheusJsonResponse::consume_stream_to_data( + RecordBatches::try_new(schema.clone(), vec![first.clone(), second.clone()]) + .unwrap() + .as_stream(), + Some("mixed_metric".to_string()), + ValueType::Vector, + ) + .await + .unwrap(); + let eager = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, vec![first, second]).unwrap(), + Some("mixed_metric".to_string()), + ValueType::Vector, + ) + .unwrap(); + + assert_eq!( + serde_json::to_value(&streamed).unwrap(), + serde_json::to_value(&eager).unwrap() + ); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = streamed + else { + panic!("expected vector response"); + }; + assert_eq!(series.len(), 2); + let histogram = series + .iter() + .find(|series| series.metric["kind"] == "histogram") + .unwrap(); + assert!(histogram.histogram.is_some()); + assert!(histogram.value.is_none()); + } + + /// A stream that lazily creates one batch per poll and asserts the previous + /// batch has been released before yielding the next one. + /// + /// Unlike `RecordBatches::as_stream`, this stream does not retain yielded + /// batches, so it can verify the consumer drops each batch before polling + /// the next one. + struct LazyDropCheckingStream { + schema: SchemaRef, + batch_index: usize, + /// Weak reference to the Arc wrapping the previously yielded batch. + /// The consumer must have dropped its clone by the next poll. + last_batch: Option>, + } + + impl LazyDropCheckingStream { + fn new(schema: SchemaRef) -> Self { + Self { + schema, + batch_index: 0, + last_batch: None, + } + } + + fn make_batch(&self) -> RecordBatch { + let base = self.batch_index as i64 * 1000; + RecordBatch::new( + self.schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values([base])) as _, + Arc::new(StringVector::from(vec![Some("host")])) as _, + Arc::new(Float64Vector::from_values([base as f64])) as _, + ], + ) + .unwrap() + } + } + + impl futures::Stream for LazyDropCheckingStream { + type Item = common_recordbatch::error::Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + // The consumer must have released the previous batch before polling + // for the next one. + if let Some(weak) = &self.last_batch { + assert!( + weak.upgrade().is_none(), + "previous batch was not released before polling the next one" + ); + } + + if self.batch_index >= 3 { + return std::task::Poll::Ready(None); + } + + let batch = Arc::new(self.make_batch()); + self.last_batch = Some(Arc::downgrade(&batch)); + self.batch_index += 1; + std::task::Poll::Ready(Some(Ok((*batch).clone()))) + } + } + + impl common_recordbatch::RecordBatchStream for LazyDropCheckingStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + None + } + } + + #[tokio::test] + async fn stream_drops_each_batch_before_polling_next() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let stream = LazyDropCheckingStream::new(schema); + + let response = PrometheusJsonResponse::consume_stream_to_data( + Box::pin(stream), + None, + ValueType::Matrix, + ) + .await; + + assert!(matches!( + response.unwrap(), + PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) if series.len() == 1 && series[0].values.len() == 3 + )); + } + #[tokio::test] async fn query_response_preserves_and_merges_promql_annotations() { let query_id = "query_response_preserves_and_merges_promql_annotations"; From 233e23b3524c9d268b7929ee2f730d48ba50f227 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 21 Sep 2026 07:09:31 +0000 Subject: [PATCH 36/63] fix(promql): keep value-field grouping labels out of matching filter propagation (#9242) * fix(promql): keep value-field grouping labels out of matching filter propagation #9202 propagates matching-label matchers to the scanned selector using the planned operands' tag columns. For an aggregated operand those are its grouping labels, and agg_modifier_to_col resolves by(...) names against the input schema only, so a value field named in by(...) is reported as a tag. A value field varies between the samples of one series, so lowering a matcher on it below sample selection (PromInstantManipulate) can drop the newest sample, promote a stale one from the lookback window, and fabricate a match the un-rewritten query does not produce. Track the by(...) labels that name value fields of the aggregated operand's input in PromPlannerContext::aggregation_field_labels, and refuse to propagate matchers on them while still requiring a matcher to name a grouping label to cross an aggregation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: sync matching_filter result fixture comment with the PR number Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/planner.rs | 151 +++++++++++++++ .../src/promql/planner/matching_filters.rs | 179 ++++++++++++++++-- .../common/tql/matching_filter.result | 53 ++++++ .../standalone/common/tql/matching_filter.sql | 36 ++++ 4 files changed, 405 insertions(+), 14 deletions(-) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index bea2cae586c..94081c2d284 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -174,6 +174,15 @@ struct PromPlannerContext { time_index_column: Option, field_columns: Vec, tag_columns: Vec, + /// `by(...)` labels of the aggregation that produced this operand that are not series tags of + /// its input, i.e. value fields (or a label an inner aggregation already reported as one). + /// + /// An aggregation reports every `by(...)` label it finds in its input schema among its tag + /// columns, and a value field named there is a group key of the aggregate rather than a + /// property of a series: its value varies between the samples of one series. Lowering a + /// matcher on it into the scan would change which samples are selected (#9242), so + /// [`matching_filters`] refuses to propagate such a matcher. + aggregation_field_labels: Vec, /// Use metric engine internal series identifier column (`__tsid`) as series key. /// /// This is enabled only when the underlying scan can provide `__tsid` (`UInt64`). The planner @@ -1359,6 +1368,8 @@ impl PromPlanner { // lhs is a literal, rhs is a column (Some(mut expr), None) => { let input = self.prom_expr_to_plan(rhs, query_engine_state).await?; + // Arithmetic against a literal preserves the series labels, so + // `aggregation_field_labels` passes through with `tag_columns` unchanged. // check if the literal is a special time expr if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) { expr = time_expr @@ -1515,6 +1526,8 @@ impl PromPlanner { binary_expr, &left_context.tag_columns, &right_context.tag_columns, + &left_context.aggregation_field_labels, + &right_context.aggregation_field_labels, ) { // A copied matcher belongs to the scan, not to the operand's identity: // `absent()` turns `selector_matcher` into the labels it reports, so the @@ -2207,6 +2220,7 @@ impl PromPlanner { self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); self.ctx.reset_table_name_and_schema(); self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()]; LogicalPlan::Extension(Extension { node: Arc::new( @@ -2599,11 +2613,27 @@ impl PromPlanner { None => { if update_ctx { self.ctx.tag_columns.clear(); + self.ctx.aggregation_field_labels.clear(); } Ok(vec![self.create_time_index_column_expr()?]) } Some(LabelModifier::Include(labels)) => { if update_ctx { + // A `by(...)` label can name a value field of the input instead of a tag. The + // aggregate still reports it among its tag columns below, but unlike a tag it + // is a group key rather than a property of a series: its value varies between + // the samples of one series, so a matcher on it must stay above sample + // selection (#9242). Record it, before the tag columns are overwritten. + self.ctx.aggregation_field_labels = labels + .labels + .iter() + .filter(|label| { + self.ctx.field_columns.contains(label) + || !self.ctx.tag_columns.contains(label) + || self.ctx.aggregation_field_labels.contains(label) + }) + .cloned() + .collect(); self.ctx.tag_columns.clear(); } let mut exprs = Vec::with_capacity(labels.labels.len()); @@ -2655,6 +2685,12 @@ impl PromPlanner { if update_ctx { // change the tag columns in context self.ctx.tag_columns = all_fields.iter().map(|col| (*col).clone()).collect(); + // `without(...)` drops the value fields of the input from its grouping labels, + // so a label stays a grouping label in name only if it survives in the input + // schema (e.g. an inner aggregation that grouped by it). + self.ctx + .aggregation_field_labels + .retain(|label| all_fields.iter().any(|col| *col == label)); } // collect remaining fields and convert to col expr @@ -3348,6 +3384,8 @@ impl PromPlanner { .cloned() .collect(); self.ctx.tag_columns = tags; + // The operand is a plain selector: its tag columns are the table's. + self.ctx.aggregation_field_labels.clear(); self.ctx.use_tsid = false; @@ -3360,6 +3398,7 @@ impl PromPlanner { self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); self.ctx.reset_table_name_and_schema(); self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); self.ctx.field_columns = vec![DEFAULT_FIELD_COLUMN.to_string()]; self.ctx.use_tsid = false; @@ -5204,6 +5243,7 @@ impl PromPlanner { self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); self.ctx.reset_table_name_and_schema(); self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); self.ctx.field_columns = vec![greptime_value().to_string()]; Ok(LogicalPlan::Extension(Extension { node: Arc::new( @@ -5295,6 +5335,7 @@ impl PromPlanner { }); // scalar plan have no tag columns self.ctx.tag_columns.clear(); + self.ctx.aggregation_field_labels.clear(); self.ctx.field_columns.clear(); self.ctx .field_columns @@ -5364,6 +5405,9 @@ impl PromPlanner { ), }); + // The absent series carries the equality matchers as labels, not the input's + // tags or value fields, so the input's field grouping labels no longer apply. + self.ctx.aggregation_field_labels.clear(); Ok(absent_plan) } @@ -13366,6 +13410,92 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom plan.display_indent().to_string() } + /// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value + /// field that is neither a primary key nor a value column of the metric, so + /// `count by(status) (...)` still reports it among the aggregation's tag columns. + async fn build_test_table_provider_with_string_field( + table_tags: &[(&str, &[&str])], + ) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + columns.push(ColumnSchema::new( + "status".to_string(), + ConcreteDataType::string_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + + async fn build_matching_filter_plan_with_string_field(query: &str) -> String { + let table_provider = build_test_table_provider_with_string_field(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() + } + #[tokio::test] async fn binary_matching_label_filter_reaches_both_operands() { for query in [ @@ -13405,6 +13535,27 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn binary_value_field_matcher_is_not_copied_across_aggregations() { + // `status` varies between the samples of one series, so filtering the other operand by it + // would drop the newest sample before sample selection (#9242). + let query = r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#; + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!( + plan.matches(r#"Utf8("ready")"#).count(), + 1, + "{query}\n{plan}" + ); + } + + #[tokio::test] + async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { + // `count`, not `sum`: the string value field is not summable. + let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); + } + #[tokio::test] async fn binary_matching_label_filter_skips_unproven_expressions() { for query in [ diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs index 3ed5ee93baf..386d95e9228 100644 --- a/src/query/src/promql/planner/matching_filters.rs +++ b/src/query/src/promql/planner/matching_filters.rs @@ -62,12 +62,30 @@ const LABEL_PRESERVING_RANGE_FUNCTIONS: [&str; 20] = [ /// as well constrain a value field, which the parsed expression does not distinguish from a /// label. For an aggregated operand they are its grouping labels, which is what makes filtering /// its input equivalent to filtering its output. +/// +/// A value field is not a series identity though: its value may vary between the samples of one +/// series, so a matcher on it must not be lowered below sample selection (`PromInstantManipulate`, +/// `PromSeriesDivide`), where discarding a sample promotes an older one out of the lookback window +/// and fabricates a match. An aggregated operand's tag columns are its grouping labels rather than +/// the tags of its input, so those labels may name value fields; a matcher therefore only crosses +/// an aggregation when it names one of the grouping labels that aggregation partitions by, and +/// `left_field_labels`/`right_field_labels` rule out the grouping labels the planner found to be +/// value fields of the aggregated operand's input +/// (`PromPlannerContext::aggregation_field_labels`). pub(super) fn propagate( binary: &BinaryExpr, left_tags: &[String], right_tags: &[String], + left_field_labels: &[String], + right_field_labels: &[String], ) -> Option { - match try_propagate(binary, left_tags, right_tags) { + match try_propagate( + binary, + left_tags, + right_tags, + left_field_labels, + right_field_labels, + ) { Ok(rewritten) => Some(rewritten), Err(reason) => { common_telemetry::debug!("Matching filter not propagated ({reason}): {binary}"); @@ -81,6 +99,8 @@ fn try_propagate( binary: &BinaryExpr, left_tags: &[String], right_tags: &[String], + left_field_labels: &[String], + right_field_labels: &[String], ) -> Result { if !matches!( binary.op.id(), @@ -101,9 +121,18 @@ fn try_propagate( .as_ref() .and_then(|modifier| modifier.matching.as_ref()); + // The grouping labels of every aggregation the operands cross, collected before `rewritten` + // takes the mutable borrows that rule out borrowing `binary` again. + let left_grouped = modifier_label_names(&binary.lhs); + let left_ignored = modifier_excluded_label_names(&binary.lhs); + let right_grouped = modifier_label_names(&binary.rhs); + let right_ignored = modifier_excluded_label_names(&binary.rhs); + let mut rewritten = binary.clone(); - let left = selector_matchers(&mut rewritten.lhs).ok_or("left operand is not a selector")?; - let right = selector_matchers(&mut rewritten.rhs).ok_or("right operand is not a selector")?; + let (left, left_restricted) = + targeted_selector(&mut rewritten.lhs).ok_or("left operand is not a selector")?; + let (right, right_restricted) = + targeted_selector(&mut rewritten.rhs).ok_or("right operand is not a selector")?; if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { return Err("selector has an or matcher group"); } @@ -113,11 +142,16 @@ fn try_propagate( !name.starts_with("__") && left_tags.contains(name) && right_tags.contains(name) + // A grouping label that is a value field of the operand's input is not a series tag. + && !left_field_labels.contains(name) + && !right_field_labels.contains(name) && match matching { None => true, Some(LabelModifier::Include(on)) => on.labels.contains(name), Some(LabelModifier::Exclude(ignoring)) => !ignoring.labels.contains(name), } + && (!left_restricted || crosses_aggregation(name, &left_grouped, &left_ignored)) + && (!right_restricted || crosses_aggregation(name, &right_grouped, &right_ignored)) }; let constraints = left .matchers @@ -150,9 +184,9 @@ fn matches_every_value(matcher: &Matcher) -> bool { /// by a grouping label drops exactly the matching output series and leaves the remaining /// aggregated values untouched. /// -/// `topk`, `bottomk` and `limitk` rank across a group and carry the input labels through, so -/// filtering before them changes the candidate set. `count_values` adds an output label that -/// does not exist in its input. +/// `topk` and `bottomk` rank across a group and carry the input labels through, so filtering +/// before them changes the candidate set. `count_values` adds an output label that does not exist +/// in its input. fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { matches!( aggregate.op.id(), @@ -168,20 +202,22 @@ fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { ) } -/// Returns the matchers of the vector selector an operand scans, or `None` when the operand's -/// output labels are not proven to be that selector's labels. -fn selector_matchers(expr: &mut Expr) -> Option<&mut Matchers> { +/// Returns the matchers of the vector selector an operand scans, together with whether the +/// operand is restricted to the grouping labels an aggregation on its path partitions by, or +/// `None` when the operand's output labels are not proven to be that selector's labels. +fn targeted_selector(expr: &mut Expr) -> Option<(&mut Matchers, bool)> { match expr { - Expr::VectorSelector(selector) => Some(&mut selector.matchers), - Expr::Paren(paren) => selector_matchers(&mut paren.expr), + Expr::VectorSelector(selector) => Some((&mut selector.matchers, false)), + Expr::Paren(paren) => targeted_selector(&mut paren.expr), Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { - selector_matchers(&mut aggregate.expr) + let (matchers, _) = targeted_selector(&mut aggregate.expr)?; + Some((matchers, aggregate.modifier.is_some())) } Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => { // Every other argument is a scalar, so position does not matter. let matrix = single_matrix_argument(&call.args.args)?; match call.args.args[matrix].as_mut() { - Expr::MatrixSelector(selector) => Some(&mut selector.vs.matchers), + Expr::MatrixSelector(selector) => Some((&mut selector.vs.matchers, false)), _ => None, } } @@ -190,6 +226,49 @@ fn selector_matchers(expr: &mut Expr) -> Option<&mut Matchers> { } } +/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the +/// scanned selector groups by (`by(...)`, the `LabelModifier::Include` of an aggregation). +fn modifier_label_names(expr: &Expr) -> Vec<&String> { + let mut grouped = Vec::new(); + collect_modifier_label_names(expr, true, &mut grouped); + grouped +} + +/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the +/// scanned selector excludes from its grouping labels (`without(...)`, the +/// `LabelModifier::Exclude` of an aggregation). +fn modifier_excluded_label_names(expr: &Expr) -> Vec<&String> { + let mut ignored = Vec::new(); + collect_modifier_label_names(expr, false, &mut ignored); + ignored +} + +fn collect_modifier_label_names<'a>(expr: &'a Expr, included: bool, out: &mut Vec<&'a String>) { + match expr { + Expr::Paren(paren) => collect_modifier_label_names(&paren.expr, included, out), + Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { + match aggregate.modifier.as_ref() { + Some(LabelModifier::Include(labels)) if included => out.extend(&labels.labels), + Some(LabelModifier::Exclude(labels)) if !included => out.extend(&labels.labels), + _ => (), + } + collect_modifier_label_names(&aggregate.expr, included, out); + } + _ => (), + } +} + +/// Whether an operand an aggregation partitions by `grouped` labels (with `without(...)`, by +/// every label but `ignored`) still carries `name` as an output label, so that a matcher on +/// `name` reaches it on both sides of that aggregation. +fn crosses_aggregation(name: &String, grouped: &[&String], ignored: &[&String]) -> bool { + if grouped.is_empty() { + !ignored.contains(&name) + } else { + grouped.contains(&name) + } +} + fn single_matrix_argument(args: &[Box]) -> Option { let mut found = None; for (index, arg) in args.iter().enumerate() { @@ -210,11 +289,30 @@ mod tests { use super::*; fn rewrite_with(query: &str, left_tags: &[&str], right_tags: &[&str]) -> Option { + rewrite_labels_with(query, left_tags, &[], right_tags, &[]) + } + + /// [`rewrite_with`] with the grouping labels of each operand that the planner found to be + /// value fields of its input rather than tags (the metric's tags are `["host", "zone"]`). + fn rewrite_labels_with( + query: &str, + left_tags: &[&str], + left_field_labels: &[&str], + right_tags: &[&str], + right_field_labels: &[&str], + ) -> Option { let Expr::Binary(binary) = parse(query).unwrap() else { panic!("expected binary") }; let owned = |tags: &[&str]| tags.iter().map(|tag| tag.to_string()).collect::>(); - propagate(&binary, &owned(left_tags), &owned(right_tags)).map(Expr::Binary) + propagate( + &binary, + &owned(left_tags), + &owned(right_tags), + &owned(left_field_labels), + &owned(right_field_labels), + ) + .map(Expr::Binary) } fn rewrite(query: &str) -> Option { @@ -298,6 +396,59 @@ mod tests { ); } + #[test] + fn does_not_propagate_grouping_labels_that_are_value_fields() { + // `status` is a value field of both metrics, though `count by(status)` reports it as a + // grouping label of the aggregate; its value varies between the samples of one series. + assert!( + rewrite_labels_with( + r#"count by(status) (a) / on(status) count by(status) (b{status="ready"})"#, + &["status"], + &["status"], + &["status"], + &["status"], + ) + .is_none() + ); + // `host` is a tag, `status` is a value field: only the tag may propagate. + assert!( + rewrite_labels_with( + r#"sum by(host, status) (a) / on(status) sum by(host, status) (b{status="ready"})"#, + &["host", "status"], + &["status"], + &["host", "status"], + &["status"], + ) + .is_none() + ); + } + + #[test] + fn propagates_grouping_labels_of_aggregated_operands() { + assert_rewrite( + r#"count by(host) (a) / on(host) count by(host) (b{host="x"})"#, + r#"count by(host) (a{host="x"}) / on(host) count by(host) (b{host="x"})"#, + ); + assert_rewrite( + r#"sum by(host) (rate(a[5m])) / on(host) sum by(host) (b{host="x"})"#, + r#"sum by(host) (rate(a{host="x"}[5m])) / on(host) sum by(host) (b{host="x"})"#, + ); + } + + #[test] + fn keeps_matchers_outside_the_grouping_labels_on_their_own_operand() { + assert_rewrite( + r#"sum by(host) (a) / on(host) sum by(host) (b{host="x",status="ready"})"#, + r#"sum by(host) (a{host="x"}) / on(host) sum by(host) (b{host="x",status="ready"})"#, + ); + // `without(host)` leaves `host` out of the grouping labels, so nothing may cross it. + assert!(rewrite(r#"avg without(host) (a) / on(host) b{host="x"}"#).is_none()); + assert_rewrite( + r#"avg without(zone) (a) / on(host) b{host="x"}"#, + r#"avg without(zone) (a{host="x"}) / on(host) b{host="x"}"#, + ); + } + #[test] fn leaves_selecting_aggregations_alone() { for query in [ diff --git a/tests/cases/standalone/common/tql/matching_filter.result b/tests/cases/standalone/common/tql/matching_filter.result index 98834ad0902..8f5bebbe8c3 100644 --- a/tests/cases/standalone/common/tql/matching_filter.result +++ b/tests/cases/standalone/common/tql/matching_filter.result @@ -332,6 +332,59 @@ tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device= | host2 | 1970-01-01T00:00:05 | 20.0 | +-------+---------------------+---------------------------------------------------------+ +-- A grouping label that is a value field, not a tag: its value varies between samples of one +-- series, so a matcher on it must not filter the scan before sample selection (#9242). +create table pr9202_a ( + ts timestamp time index, + host string, + val double, + `status` string, + primary key (host) +); + +Affected Rows: 0 + +create table pr9202_b ( + ts timestamp time index, + host string, + val double, + `status` string, + primary key (host) +); + +Affected Rows: 0 + +insert into pr9202_a values + (0, 'h1', 10, 'ready'), + (5000, 'h1', 20, 'busy'); + +Affected Rows: 2 + +insert into pr9202_b values + (5000, 'h1', 5, 'ready'); + +Affected Rows: 1 + +-- The 5s sample of pr9202_a has status="busy" and no partner; filtering it before sample +-- selection would fall back to the stale "ready" sample and fabricate a result row. +tql eval (5, 5, '1s', '1m') + count by(status) (pr9202_a) + / on(status) + count by(status) (pr9202_b{status="ready"}); + ++--------+----+-------------------------------------------------------------+-------------------------------------------------------------------+ +| status | ts | pr9202_a.count(pr9202_a.val) / pr9202_b.count(pr9202_b.val) | pr9202_a.count(pr9202_a.status) / pr9202_b.count(pr9202_b.status) | ++--------+----+-------------------------------------------------------------+-------------------------------------------------------------------+ ++--------+----+-------------------------------------------------------------+-------------------------------------------------------------------+ + +drop table pr9202_a; + +Affected Rows: 0 + +drop table pr9202_b; + +Affected Rows: 0 + -- `topk` preserves its input labels, and filtering before `topk` changes the set of -- candidates it ranks, so the rewrite must not reach its input. -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/tql/matching_filter.sql b/tests/cases/standalone/common/tql/matching_filter.sql index fa2bc3371d5..72da3da70a0 100644 --- a/tests/cases/standalone/common/tql/matching_filter.sql +++ b/tests/cases/standalone/common/tql/matching_filter.sql @@ -115,6 +115,42 @@ tql eval(0, 5, '5s') sum by(host)(counter_metric{device="eth0"}) / on(host) sum -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device="eth0"}); +-- A grouping label that is a value field, not a tag: its value varies between samples of one +-- series, so a matcher on it must not filter the scan before sample selection (#9242). +create table pr9202_a ( + ts timestamp time index, + host string, + val double, + `status` string, + primary key (host) +); + +create table pr9202_b ( + ts timestamp time index, + host string, + val double, + `status` string, + primary key (host) +); + +insert into pr9202_a values + (0, 'h1', 10, 'ready'), + (5000, 'h1', 20, 'busy'); + +insert into pr9202_b values + (5000, 'h1', 5, 'ready'); + +-- The 5s sample of pr9202_a has status="busy" and no partner; filtering it before sample +-- selection would fall back to the stale "ready" sample and fabricate a result row. +tql eval (5, 5, '1s', '1m') + count by(status) (pr9202_a) + / on(status) + count by(status) (pr9202_b{status="ready"}); + +drop table pr9202_a; + +drop table pr9202_b; + -- `topk` preserves its input labels, and filtering before `topk` changes the set of -- candidates it ranks, so the rewrite must not reach its input. -- SQLNESS SORT_RESULT 3 1 From b332bb67d69f552ba40230d301e5857bd895c682 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 21 Sep 2026 07:19:20 +0000 Subject: [PATCH 37/63] perf(telemetry): remove the shared lock from TraceLayer (#9269) * perf(telemetry): remove the shared lock from TraceLayer Signed-off-by: Dennis Zhuang * fix(telemetry): stop collecting span data while tracing is disabled Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.lock | 1 - src/common/telemetry/src/logging.rs | 315 +++++++++++++++++++++++++--- src/servers/Cargo.toml | 1 - src/servers/src/http/dyn_trace.rs | 16 +- 4 files changed, 286 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7ebd6bdb3db..ba3f508486e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13681,7 +13681,6 @@ dependencies = [ "tower 0.5.2", "tower-http 0.6.11", "tracing", - "tracing-opentelemetry", "urlencoding", "uuid", "vrl", diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index 83ceb7aef9d..a9b50930786 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -18,7 +18,8 @@ mod file_retention; use std::collections::HashMap; use std::env; use std::io::IsTerminal; -use std::sync::{Arc, Mutex, Once, RwLock}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, Once}; use std::time::Duration; use common_base::readable_size::ReadableSize; @@ -60,35 +61,55 @@ pub static LOG_RELOAD_HANDLE: OnceCell, Registry>; type OtelTraceLayer = tracing_opentelemetry::OpenTelemetryLayer; +struct TraceLayerState { + enabled: AtomicBool, + layer: OnceCell, +} + #[derive(Clone)] pub struct TraceReloadHandle { - inner: Arc>>, + inner: Arc, } impl TraceReloadHandle { - fn new(inner: Arc>>) -> Self { + fn new(inner: Arc) -> Self { Self { inner } } - pub fn reload(&self, new_layer: Option) { - let mut guard = self.inner.write().unwrap(); - *guard = new_layer; - drop(guard); + /// Enables or disables OTLP data collection, initializing it on first enable. + /// Disabling stops new spans, events, fields and links. Existing spans still + /// finish their lifecycle and export on close. + pub fn set_enabled(&self, enabled: bool) -> Result<(), &'static str> { + self.set_enabled_with(enabled, || { + get_or_init_tracer().map(|tracer| tracing_opentelemetry::layer().with_tracer(tracer)) + }) + } + fn set_enabled_with( + &self, + enabled: bool, + init: impl FnOnce() -> Result, + ) -> Result<(), &'static str> { + if enabled { + self.inner.layer.get_or_try_init(init)?; + } + self.inner.enabled.store(enabled, Ordering::Release); callsite::rebuild_interest_cache(); + Ok(()) } } -/// A tracing layer that can be dynamically reloaded. -/// -/// Mostly copied from [`tracing_subscriber::reload::Layer`]. +/// An OTLP layer with a runtime switch and a stable address for downcasts. struct TraceLayer { - inner: Arc>>, + inner: Arc, } impl TraceLayer { fn new(initial: Option) -> (Self, TraceReloadHandle) { - let inner = Arc::new(RwLock::new(initial)); + let inner = Arc::new(TraceLayerState { + enabled: AtomicBool::new(initial.is_some()), + layer: initial.map(OnceCell::with_value).unwrap_or_default(), + }); ( Self { inner: inner.clone(), @@ -98,17 +119,11 @@ impl TraceLayer { } fn with_layer(&self, f: impl FnOnce(&OtelTraceLayer) -> R) -> Option { - self.inner - .read() - .ok() - .and_then(|guard| guard.as_ref().map(f)) + self.inner.layer.get().map(f) } - fn with_layer_mut(&self, f: impl FnOnce(&mut OtelTraceLayer) -> R) -> Option { - self.inner - .write() - .ok() - .and_then(|mut guard| guard.as_mut().map(f)) + fn is_enabled(&self) -> bool { + self.inner.enabled.load(Ordering::Acquire) } } @@ -117,10 +132,6 @@ impl tracing_subscriber::Layer for TraceLayer { let _ = self.with_layer(|layer| layer.on_register_dispatch(subscriber)); } - fn on_layer(&mut self, subscriber: &mut DynSubscriber) { - let _ = self.with_layer_mut(|layer| layer.on_layer(subscriber)); - } - fn register_callsite( &self, metadata: &'static tracing::Metadata<'static>, @@ -144,7 +155,9 @@ impl tracing_subscriber::Layer for TraceLayer { id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, DynSubscriber>, ) { - let _ = self.with_layer(|layer| layer.on_new_span(attrs, id, ctx)); + if self.is_enabled() { + let _ = self.with_layer(|layer| layer.on_new_span(attrs, id, ctx)); + } } fn max_level_hint(&self) -> Option { @@ -157,7 +170,9 @@ impl tracing_subscriber::Layer for TraceLayer { values: &tracing::span::Record<'_>, ctx: tracing_subscriber::layer::Context<'_, DynSubscriber>, ) { - let _ = self.with_layer(|layer| layer.on_record(span, values, ctx)); + if self.is_enabled() { + let _ = self.with_layer(|layer| layer.on_record(span, values, ctx)); + } } fn on_follows_from( @@ -166,7 +181,9 @@ impl tracing_subscriber::Layer for TraceLayer { follows: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, DynSubscriber>, ) { - let _ = self.with_layer(|layer| layer.on_follows_from(span, follows, ctx)); + if self.is_enabled() { + let _ = self.with_layer(|layer| layer.on_follows_from(span, follows, ctx)); + } } fn event_enabled( @@ -183,7 +200,9 @@ impl tracing_subscriber::Layer for TraceLayer { event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, DynSubscriber>, ) { - let _ = self.with_layer(|layer| layer.on_event(event, ctx)); + if self.is_enabled() { + let _ = self.with_layer(|layer| layer.on_event(event, ctx)); + } } fn on_enter( @@ -220,11 +239,13 @@ impl tracing_subscriber::Layer for TraceLayer { } unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> { - self.inner.read().ok().and_then(|guard| { - guard - .as_ref() - .and_then(|layer| unsafe { layer.downcast_raw(id) }) - }) + // Keep downcasts available while disabled: an in-flight WithContext + // callback may still need the layer. OnceCell keeps both addresses valid + // for the subscriber's lifetime, even across concurrent toggles. + self.inner + .layer + .get() + .and_then(|layer| unsafe { layer.downcast_raw(id) }) } } @@ -764,8 +785,234 @@ where #[cfg(test)] mod tests { + use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + + use opentelemetry::trace::TraceContextExt; + use opentelemetry_sdk::trace::SdkTracerProvider; + use tracing_opentelemetry::OpenTelemetrySpanExt; + use super::*; + #[test] + fn test_trace_switch_preserves_active_spans() { + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .build(); + for initially_enabled in [false, true] { + let new_layer = + || tracing_opentelemetry::layer().with_tracer(provider.tracer("switch")); + let (filter, _) = tracing_subscriber::reload::Layer::new( + Targets::new().with_default(tracing::Level::INFO), + ); + let (layer, handle) = TraceLayer::new(initially_enabled.then(new_layer)); + let dispatch = tracing::Dispatch::new(Registry::default().with(filter).with(layer)); + tracing::dispatcher::with_default(&dispatch, || { + let initial = tracing::info_span!("initial"); + assert!(!initial.is_disabled()); + assert_eq!( + initial.context().span().span_context().is_valid(), + initially_enabled + ); + handle.set_enabled_with(true, || Ok(new_layer())).unwrap(); + let parent = tracing::info_span!("parent"); + let parent_context = parent.context(); + assert!(parent_context.span().span_context().is_valid()); + let previous_context = opentelemetry::Context::current(); + let entered = parent.enter(); + handle.set_enabled(false).unwrap(); + let disabled = tracing::info_span!("disabled"); + assert!(!disabled.is_disabled()); + assert!(!disabled.context().span().span_context().is_valid()); + assert_eq!( + parent.context().span().span_context(), + parent_context.span().span_context() + ); + assert!(dispatch.downcast_ref::().is_some()); + drop(entered); + assert_eq!( + opentelemetry::Context::current().span().span_context(), + previous_context.span().span_context() + ); + let disabled_entered = disabled.enter(); + handle + .set_enabled_with(true, || panic!("layer must not be replaced")) + .unwrap(); + drop(disabled_entered); + let child = tracing::info_span!(parent: &parent, "child"); + assert_eq!( + child.context().span().span_context().trace_id(), + parent_context.span().span_context().trace_id() + ); + assert!(!disabled.context().span().span_context().is_valid()); + }); + } + } + + #[test] + fn test_trace_switch_retries_initialization() { + let (filter, _) = tracing_subscriber::reload::Layer::new( + Targets::new().with_default(tracing::Level::INFO), + ); + let (layer, handle) = TraceLayer::new(None); + let subscriber = Registry::default().with(filter).with(layer); + tracing::subscriber::with_default(subscriber, || { + handle + .set_enabled_with(false, || panic!("disabling must not initialize OTLP")) + .unwrap(); + assert_eq!( + handle.set_enabled_with(true, || Err("initialization failed")), + Err("initialization failed") + ); + let failed = tracing::info_span!("after_failed_enable"); + assert!(!failed.context().span().span_context().is_valid()); + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .build(); + handle + .set_enabled_with(true, || { + Ok(tracing_opentelemetry::layer().with_tracer(provider.tracer("retry"))) + }) + .unwrap(); + let enabled = tracing::info_span!("after_successful_enable"); + assert!(enabled.context().span().span_context().is_valid()); + }); + } + + #[test] + fn test_trace_switch_stops_collecting_fields_when_disabled() { + struct CountFormatting(AtomicUsize); + + impl std::fmt::Debug for CountFormatting { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fetch_add(1, Ordering::Relaxed); + f.write_str("value") + } + } + + let value = CountFormatting(AtomicUsize::new(0)); + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .build(); + let (filter, _) = tracing_subscriber::reload::Layer::new( + Targets::new().with_default(tracing::Level::INFO), + ); + let (layer, handle) = TraceLayer::new(Some( + tracing_opentelemetry::layer().with_tracer(provider.tracer("fields")), + )); + tracing::subscriber::with_default(Registry::default().with(filter).with(layer), || { + let admitted = tracing::info_span!("admitted", field = tracing::field::Empty); + admitted.record("field", tracing::field::debug(&value)); + admitted.in_scope(|| tracing::info!(value = ?value)); + assert_eq!(value.0.load(Ordering::Relaxed), 2); + + handle.set_enabled(false).unwrap(); + let unadmitted = tracing::info_span!("unadmitted", field = tracing::field::Empty); + for span in [&admitted, &unadmitted] { + span.record("field", tracing::field::debug(&value)); + span.in_scope(|| tracing::info!(value = ?value)); + } + assert_eq!(value.0.load(Ordering::Relaxed), 2); + + handle + .set_enabled_with(true, || panic!("layer must not be replaced")) + .unwrap(); + admitted.record("field", tracing::field::debug(&value)); + admitted.in_scope(|| tracing::info!(value = ?value)); + assert_eq!(value.0.load(Ordering::Relaxed), 4); + }); + } + + #[test] + fn test_trace_switch_finishes_admitted_spans() { + #[derive(Debug)] + struct Exporter(std::sync::mpsc::Sender); + + impl opentelemetry_sdk::trace::SpanExporter for Exporter { + async fn export( + &self, + batch: Vec, + ) -> opentelemetry_sdk::error::OTelSdkResult { + for span in batch { + self.0.send(span.name.into_owned()).unwrap(); + } + Ok(()) + } + } + + let (sender, receiver) = std::sync::mpsc::channel(); + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .with_simple_exporter(Exporter(sender)) + .build(); + let (filter, _) = tracing_subscriber::reload::Layer::new( + Targets::new().with_default(tracing::Level::INFO), + ); + let (layer, handle) = TraceLayer::new(Some( + tracing_opentelemetry::layer().with_tracer(provider.tracer("export")), + )); + tracing::subscriber::with_default(Registry::default().with(filter).with(layer), || { + let active = tracing::info_span!("admitted"); + handle.set_enabled(false).unwrap(); + let disabled = tracing::info_span!("not_admitted"); + drop(active); + drop(disabled); + }); + provider.force_flush().unwrap(); + assert_eq!(receiver.try_iter().collect::>(), ["admitted"]); + } + + #[test] + fn test_trace_switch_concurrent_context_access() { + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::AlwaysOn) + .build(); + let (filter, _) = tracing_subscriber::reload::Layer::new( + Targets::new().with_default(tracing::Level::INFO), + ); + let (layer, handle) = TraceLayer::new(Some( + tracing_opentelemetry::layer().with_tracer(provider.tracer("concurrent")), + )); + let dispatch = tracing::Dispatch::new(Registry::default().with(filter).with(layer)); + let parent = tracing::dispatcher::with_default(&dispatch, || { + tracing::info_span!("concurrent_parent") + }); + let context = parent.context(); + let barrier = Barrier::new(4); + std::thread::scope(|scope| { + for _ in 0..3 { + scope.spawn(|| { + tracing::dispatcher::with_default(&dispatch, || { + let _entered = parent.enter(); + barrier.wait(); + for _ in 0..500 { + assert_eq!( + parent.context().span().span_context(), + context.span().span_context() + ); + { + let child = tracing::info_span!("concurrent_child"); + let _entered = child.enter(); + let _ = child.context(); + tracing::info!("concurrent event"); + } + assert_eq!( + opentelemetry::Context::current().span().span_context(), + context.span().span_context() + ); + } + }); + }); + } + barrier.wait(); + for i in 0..1000 { + handle + .set_enabled_with(i % 2 == 0, || panic!("layer must not be replaced")) + .unwrap(); + } + }); + } + #[test] fn test_logging_options_deserialization_default() { let json = r#"{}"#; diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 5fa10313e04..8ffd97e2c0f 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -140,7 +140,6 @@ tonic-web = "0.14" tower = { workspace = true, features = ["full"] } tower-http = { version = "0.6", features = ["full"] } tracing.workspace = true -tracing-opentelemetry.workspace = true urlencoding = "2.1" uuid.workspace = true vrl.workspace = true diff --git a/src/servers/src/http/dyn_trace.rs b/src/servers/src/http/dyn_trace.rs index dcdb74c56ab..7fe1880cf6b 100644 --- a/src/servers/src/http/dyn_trace.rs +++ b/src/servers/src/http/dyn_trace.rs @@ -14,7 +14,7 @@ use axum::http::StatusCode; use axum::response::IntoResponse; -use common_telemetry::{TRACE_RELOAD_HANDLE, get_or_init_tracer, info}; +use common_telemetry::{TRACE_RELOAD_HANDLE, info}; use crate::error::{InvalidParameterSnafu, Result}; @@ -34,20 +34,14 @@ pub async fn dyn_trace_handler(enable_str: String) -> Result )); }; - if enable { - let tracer = match get_or_init_tracer() { - Ok(tracer) => tracer, - Err(reason) => { - return Ok((StatusCode::SERVICE_UNAVAILABLE, reason.to_string())); - } - }; + if let Err(reason) = trace_reload_handle.set_enabled(enable) { + return Ok((StatusCode::SERVICE_UNAVAILABLE, reason.to_string())); + } - let trace_layer = tracing_opentelemetry::layer().with_tracer(tracer); - trace_reload_handle.reload(Some(trace_layer)); + if enable { info!("trace enabled"); Ok((StatusCode::OK, "trace enabled".to_string())) } else { - trace_reload_handle.reload(None); info!("trace disabled"); Ok((StatusCode::OK, "trace disabled".to_string())) } From 66d38e8e1c007d17ffec1fae0e3b31cef960e305 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:27:22 +0000 Subject: [PATCH 38/63] feat: add database ingestion admission through metering (#9239) * feat: add `ingest_rows_rate_limit` database option Signed-off-by: Lei, HUANG * feat: add `InsertLimitInterceptor` hook to `Inserter` Signed-off-by: Lei, HUANG * fix: attribute insert limit checks to the target table's database Signed-off-by: Lei, HUANG * refactor: unify write admission through metering Signed-off-by: shuiyisong * fix: enforce write admission for pending row batches Signed-off-by: shuiyisong * feat: distinguish internal requests for ingestion metering Signed-off-by: shuiyisong * fix: admit split ingestion requests once per database Signed-off-by: shuiyisong * fix: OpenTSDB throws error reason Signed-off-by: shuiyisong * chore: use meter crate main rev Signed-off-by: shuiyisong * fix: exclude database ingest rate limit from table options Signed-off-by: shuiyisong * refactor: reserve channel 255 for internal requests Signed-off-by: shuiyisong --------- Signed-off-by: Lei, HUANG Signed-off-by: shuiyisong Co-authored-by: Lei, HUANG --- Cargo.lock | 6 +- Cargo.toml | 4 +- src/common/base/src/protocol.rs | 7 +- src/flow/AGENTS.md | 2 + src/flow/src/batching_mode/frontend_client.rs | 8 +- src/frontend/AGENTS.md | 11 +- src/frontend/src/instance/grpc.rs | 1 + src/frontend/src/instance/log_handler.rs | 5 + src/frontend/src/instance/otlp.rs | 17 +- src/frontend/src/instance/otlp/README.md | 4 + src/frontend/src/instance/prom_store.rs | 15 +- src/frontend/src/server.rs | 15 +- src/operator/AGENTS.md | 8 + src/operator/src/bulk_insert.rs | 17 +- src/operator/src/error.rs | 11 +- src/operator/src/insert.rs | 437 +++++++++++++++++- src/operator/src/statement/ddl.rs | 7 +- src/servers/AGENTS.md | 6 + src/servers/Cargo.toml | 2 + src/servers/src/batcher/logical_table.rs | 17 + src/servers/src/error.rs | 12 +- src/servers/src/grpc/context_auth.rs | 39 +- src/servers/src/grpc/database.rs | 17 +- src/servers/src/grpc/flight.rs | 10 +- src/servers/src/grpc/greptime_handler.rs | 3 +- src/servers/src/grpc/prom_query_gateway.rs | 13 +- src/servers/src/http/opentsdb.rs | 16 +- src/servers/src/http/prom_store.rs | 12 +- src/servers/src/otel_arrow.rs | 9 +- src/servers/tests/http/opentsdb_test.rs | 14 + src/session/src/context.rs | 48 +- src/sql/src/parsers/create_parser.rs | 12 + src/table/src/requests.rs | 33 ++ tests-integration/src/tests/instance_test.rs | 51 ++ 34 files changed, 827 insertions(+), 62 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ba3f508486e..de3bc7ca765 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8556,7 +8556,7 @@ dependencies = [ [[package]] name = "meter-core" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-meter.git?rev=5618e779cf2bb4755b499c630fba4c35e91898cb#5618e779cf2bb4755b499c630fba4c35e91898cb" +source = "git+https://github.com/GreptimeTeam/greptime-meter.git?rev=f743c3a5fe8c1f57363eac5866c356b7d91f9b39#f743c3a5fe8c1f57363eac5866c356b7d91f9b39" dependencies = [ "anymap2", "once_cell", @@ -8567,7 +8567,7 @@ dependencies = [ [[package]] name = "meter-macros" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-meter.git?rev=5618e779cf2bb4755b499c630fba4c35e91898cb#5618e779cf2bb4755b499c630fba4c35e91898cb" +source = "git+https://github.com/GreptimeTeam/greptime-meter.git?rev=f743c3a5fe8c1f57363eac5866c356b7d91f9b39#f743c3a5fe8c1f57363eac5866c356b7d91f9b39" dependencies = [ "meter-core", ] @@ -13617,6 +13617,8 @@ dependencies = [ "local-ip-address", "log-query", "loki-proto", + "meter-core", + "meter-macros", "metric-engine", "mime_guess", "moka", diff --git a/Cargo.toml b/Cargo.toml index 27e829a8cd0..fa01426bbbf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -178,7 +178,7 @@ jsonb = { git = "https://github.com/GreptimeTeam/jsonb", rev = "dd1e2341e3df35e9 lazy_static = "1.4" local-ip-address = "0.6" loki-proto = { git = "https://github.com/GreptimeTeam/loki-proto.git", rev = "f69c8924c4babe516373e26a4118be82d976629c" } -meter-core = { git = "https://github.com/GreptimeTeam/greptime-meter.git", rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" } +meter-core = { git = "https://github.com/GreptimeTeam/greptime-meter.git", rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" } mockall = "0.13" moka = "0.12" nalgebra = "0.33" @@ -357,7 +357,7 @@ table = { path = "src/table" } [workspace.dependencies.meter-macros] git = "https://github.com/GreptimeTeam/greptime-meter.git" -rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" +rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } diff --git a/src/common/base/src/protocol.rs b/src/common/base/src/protocol.rs index 66704dbc18b..f9ade9630c5 100644 --- a/src/common/base/src/protocol.rs +++ b/src/common/base/src/protocol.rs @@ -14,7 +14,7 @@ use std::fmt::{Display, Formatter}; -/// The protocol through which a query is received. +/// The protocol or internal subsystem through which a query is received. #[derive(Debug, PartialEq, Default, Clone, Copy, strum::FromRepr)] #[repr(u8)] pub enum Channel { @@ -35,6 +35,8 @@ pub enum Channel { Log = 12, Promql = 13, Splunk = 14, + /// Trusted internal requests and local subsystem execution. + Internal = 255, } impl From for Channel { @@ -70,6 +72,7 @@ impl AsRef for Channel { Self::Log => "log", Self::Promql => "promql", Self::Splunk => "splunk", + Self::Internal => "internal", } } } @@ -95,6 +98,7 @@ mod tests { (12, "log"), (13, "promql"), (14, "splunk"), + (255, "internal"), ]; for (value, name) in expected { @@ -102,5 +106,6 @@ mod tests { } assert_eq!("unknown", Channel::from(0).as_ref()); assert_eq!("unknown", Channel::from(15).as_ref()); + assert_eq!("unknown", Channel::from(256).as_ref()); } } diff --git a/src/flow/AGENTS.md b/src/flow/AGENTS.md index 0b024e96fbd..266621f508f 100644 --- a/src/flow/AGENTS.md +++ b/src/flow/AGENTS.md @@ -54,6 +54,8 @@ default, and sink-creation helpers. ## Public surface - gRPC `Flow` service in `src/flow/src/server.rs`. +- Local streaming sink writes and batching query execution use `Channel::Internal` + so meter collectors can account for sink writes without charging ingestion quota. - `FlowEngine` trait in `src/flow/src/engine.rs`. - Started from the `cmd` crate via `FlownodeBuilder` / `FlownodeInstance`. diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index efec4cdea63..aa2f2aa072a 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -34,7 +34,7 @@ use query::options::{FlowQueryExtensions, QueryOptions}; use rand::rng; use rand::seq::SliceRandom; use servers::query_handler::grpc::GrpcQueryHandler; -use session::context::{QueryContextBuilder, QueryContextRef}; +use session::context::{Channel, QueryContextBuilder, QueryContextRef}; use session::hints::READ_PREFERENCE_HINT; use snafu::{OptionExt, ResultExt}; use tokio::sync::SetOnce; @@ -521,6 +521,7 @@ impl FrontendClient { .current_catalog(catalog.to_string()) .current_schema(schema.to_string()) .extensions(extensions_map) + .channel(Channel::Internal) .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs.clone()))) .build(); let ctx = Arc::new(ctx); @@ -589,6 +590,7 @@ impl FrontendClient { let ctx = QueryContextBuilder::default() .current_catalog(catalog.to_string()) .current_schema(schema.to_string()) + .channel(Channel::Internal) .extensions(HashMap::from([( QUERY_PARALLELISM_HINT.to_string(), query.parallelism.to_string(), @@ -894,8 +896,9 @@ mod tests { async fn do_query( &self, _query: Request, - _ctx: QueryContextRef, + ctx: QueryContextRef, ) -> std::result::Result { + assert_eq!(ctx.channel(), Channel::Internal); self.calls.fetch_add(1, Ordering::SeqCst); Ok(Output::new_with_affected_rows(1)) } @@ -941,6 +944,7 @@ mod tests { ctx: QueryContextRef, ) -> std::result::Result { assert_eq!(ctx.extension("flow.return_region_seq"), Some("true")); + assert_eq!(ctx.channel(), Channel::Internal); Ok(Output::new_with_affected_rows(1)) } } diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 9740d160164..d552d73cf39 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -41,8 +41,15 @@ remote datanodes via `operator`/`client`. `StatementExecutor`. Distributed scans enter through `region_query.rs`. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional - auto-create, partition routing) → local `RegionServer` (standalone) or RPC to - datanodes (distributed). + auto-create, partition routing, meter admission) → local `RegionServer` + (standalone) or RPC to datanodes (distributed). 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`). + The returned context covers chunks and derived writes while preserving WCU + accounting and the original protocol channel. +- Internal gRPC listeners mark requests with `Channel::Internal` in middleware + (`server.rs`), including requests handled by Enterprise Flight wrappers. - **Table batching** (`instance/builder.rs`): protocol entry points opt in through `QueryContext`. The primary inserter prepares eligible ordinary-table writes diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 2d89b5489df..e2435b65689 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -454,6 +454,7 @@ impl Instance { request.record_batch, request.schema_bytes, ctx.skip_wal(), + ctx.channel(), ) .await .context(TableOperationSnafu)?; diff --git a/src/frontend/src/instance/log_handler.rs b/src/frontend/src/instance/log_handler.rs index 051018c9a55..44738b9dcc1 100644 --- a/src/frontend/src/instance/log_handler.rs +++ b/src/frontend/src/instance/log_handler.rs @@ -80,6 +80,11 @@ impl PipelineHandler for Instance { prepared.push((Arc::new(ctx.fork()), log)); } + operator::insert::admit_row_insert_batches(&mut prepared) + .await + .map_err(BoxedError::new) + .context(servers::error::ExecuteGrpcQuerySnafu)?; + let mut outputs = Vec::with_capacity(prepared.len()); for (ctx, log) in prepared { outputs.push(self.handle_log_inserts(log, ctx).await); diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index d72b79a4275..fa40c7eb59a 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -30,6 +30,7 @@ use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_telemetry::{tracing, warn}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; +use operator::insert::{admit_row_insert_batches, admit_write}; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, AuthSnafu, Result as ServerResult}; @@ -141,6 +142,10 @@ impl OpenTelemetryProtocolHandler for Instance { self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) .context(AuthSnafu)?; + let ctx = admit_write(rows as u64, &ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; OTLP_METRICS_ROWS.inc_by(rows as u64); @@ -239,6 +244,12 @@ impl OpenTelemetryProtocolHandler for Instance { let targets = trace_permission_targets(&table_name, &spans, &ctx); self.check_table_permission(&ctx, PermissionReq::Action(OTLP_WRITE), targets) .context(AuthSnafu)?; + // Count the external spans once, before chunking, retries, or derived tables. + let rows = spans.iter().map(|group| group.spans.len() as u64).sum(); + let ctx = admit_write(rows, &ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; self.ingest_trace_spans(&pipeline, table_name, spans, &conventions, ctx) .await } @@ -284,11 +295,15 @@ impl OpenTelemetryProtocolHandler for Instance { ) .await?; - let batches = opt_req.as_req_iter(ctx).collect::>(); + let mut batches = opt_req.as_req_iter(ctx).collect::>(); for (temp_ctx, requests) in &batches { self.check_row_insert_permission(requests, temp_ctx, PermissionReq::Action(OTLP_WRITE)) .context(AuthSnafu)?; } + admit_row_insert_batches(&mut batches) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; let mut outputs = Vec::with_capacity(batches.len()); for (temp_ctx, requests) in batches { diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index 0ce3e94a26a..c123f1caae1 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -68,6 +68,10 @@ The parser in [`servers/src/otlp/trace/span.rs`](../../../../servers/src/otlp/trace/span.rs) produces one `TraceSpanGroup` per resource/scope pair. +After permission checks, admission counts all spans in the request once per +database. Chunks, schema retries, and derived lookup-table writes retain that +admission; their actual writes still contribute WCU accounting. + `ingest_trace_spans` in [`trace_ingest.rs`](trace_ingest.rs) splits the spans in each group into owned chunks. `trace_ingest_chunk_size` defaults to 512; setting it to 0 disables splitting. The option is defined in diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 9aa228715c0..d972abc9112 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -527,11 +527,16 @@ impl PromStoreProtocolHandler for Instance { ) -> ServerResult>> { let mut prepared = Vec::with_capacity(requests.len()); for (ctx, request) in requests { - prepared.push(self.prepare_prom_store_write(request, ctx).await?); + let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?; + prepared.push((ctx, request)); } + operator::insert::admit_row_insert_batches(&mut prepared) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; let mut outputs = Vec::with_capacity(prepared.len()); - for (request, ctx) in prepared { + for (ctx, request) in prepared { let output = self.write_prepared(request, ctx, with_metric_engine).await; let failed = output.is_err(); outputs.push(output); @@ -765,9 +770,13 @@ impl PromStoreProtocolHandler for ExportMetricHandler { async fn write_all( &self, - requests: Vec<(QueryContextRef, RowInsertRequests)>, + mut requests: Vec<(QueryContextRef, RowInsertRequests)>, with_metric_engine: bool, ) -> ServerResult>> { + operator::insert::admit_row_insert_batches(&mut requests) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; let mut outputs = Vec::with_capacity(requests.len()); for (ctx, request) in requests { let output = self.write_prepared(request, ctx, with_metric_engine).await; diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 532b0d08b11..017bba53fba 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -45,6 +45,7 @@ use servers::postgres::PostgresServer; use servers::request_memory_limiter::ServerMemoryLimiter; use servers::server::{Server, ServerHandlers}; use servers::tls::{ReloadableTlsServerConfig, maybe_watch_server_tls_config}; +use session::context::Channel; use snafu::ResultExt; use tonic::Status; @@ -280,11 +281,15 @@ where .flight_handler(flight_handler) .add_layer(axum::middleware::from_fn_with_state( self.instance.clone(), - async move |State(state): State>, request: Request, next: Next| { + move |State(state): State>, mut request: Request, next: Next| async move { if state.is_suspended() { let status = Status::from(servers::error::SuspendedSnafu.build()); return status.into_http(); } + // The listener owns this marker; clients cannot set request extensions. + if !external { + request.extensions_mut().insert(Channel::Internal); + } next.run(request).await }, )); @@ -612,6 +617,10 @@ mod tests { &self, request: Request, ) -> std::result::Result>, Status> { + assert_eq!( + request.extensions().get::(), + Some(&Channel::Internal) + ); self.do_get_calls.fetch_add(1, Ordering::SeqCst); self.inner.do_get(request).await } @@ -620,6 +629,10 @@ mod tests { &self, request: Request>, ) -> std::result::Result>, Status> { + assert_eq!( + request.extensions().get::(), + Some(&Channel::Internal) + ); self.do_put_calls.fetch_add(1, Ordering::SeqCst); self.inner.do_put(request).await } diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 761d8e5080a..62853a18f61 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -38,3 +38,11 @@ cargo sqlness bare -t ``` Check standalone and distributed routing when changing region dispatch. + +Write admission uses fallible `write_meter!` calls in `Inserter::do_request` +and `handle_bulk_insert`, before flow side effects and datanode dispatch. Bulk +records carry row counts and the request channel with zero WCU. Validate both +the default `meter-macros/noop` build and active metering when changing this path. +Finite requests split internally use `admit_write` / `admit_row_insert_batches` +before dispatch. Their database-scoped `QueryContext` admission prevents a second +row debit while `do_request` still records WCU for the actual writes. diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 51a08d2f324..f1f57a64479 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -30,7 +30,9 @@ use common_grpc::flight::{FlightEncoder, FlightMessage, record_batch_to_ipc}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; use futures::future::{join_all, try_join_all}; -use session::context::QueryContextRef; +use meter_core::data::MeterRecord; +use meter_macros::write_meter; +use session::context::{Channel, QueryContextRef}; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; @@ -128,6 +130,7 @@ impl Inserter { record_batch: RecordBatch, schema_bytes: Bytes, skip_wal: bool, + channel: Channel, ) -> Result { let table_info = table.table_info(); let table_id = table_info.table_id(); @@ -137,6 +140,18 @@ impl Inserter { return Ok(0); } + // The zero value is WCU, not bytes. Bulk writes have no WCU accounting; + // preserve that behavior while admitting their rows before dispatch. + write_meter!(MeterRecord::new( + table_info.catalog_name.clone(), + table_info.schema_name.clone(), + 0, + record_batch.num_rows() as u64, + channel as u8, + )) + .await + .context(error::WriteRejectedSnafu)?; + let body_size = raw_flight_data.data_body.len(); // TODO(yingwen): Fill record batch impure default values. Note that we should override `raw_flight_data` if we have to fill defaults. // notify flownode to update dirty timestamps if flow is configured. diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 5a005250b94..40f968bc413 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -122,6 +122,14 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Write rejected: {error}"))] + WriteRejected { + #[snafu(source)] + error: meter_core::collect::WriteRejected, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to insert data"))] RequestInserts { #[snafu(implicit)] @@ -1066,6 +1074,7 @@ impl ErrorExt for Error { Error::AlterExprToRequest { source, .. } => source.status_code(), Error::External { source, .. } => source.status_code(), Error::BatchFlush { source, .. } => source.status_code(), + Error::WriteRejected { .. } => StatusCode::RateLimited, Error::FindTablePartitionRule { source, .. } | Error::SplitInsert { source, .. } | Error::SplitDelete { source, .. } @@ -1122,7 +1131,7 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { Error::ReadObject { error, .. } => retry_hint_from_opendal_error(error), - Error::ReadParquetMetadata { .. } => RetryHint::Retryable, + Error::ReadParquetMetadata { .. } | Error::WriteRejected { .. } => RetryHint::Retryable, Error::InvalidateTableCache { source, .. } | Error::ExecuteDdl { source, .. } | Error::RequestInserts { source, .. } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 35727ae9b35..989a4aa399d 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -48,6 +48,7 @@ use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; use datatypes::schema::SkippingIndexOptions; use futures_util::future; +use meter_core::data::MeterRecord; use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; @@ -78,6 +79,7 @@ use crate::batcher::PendingRowsBatcher; use crate::error::{ CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu, InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, + WriteRejectedSnafu, }; use crate::expr_helper; use crate::region_req_factory::RegionRequestFactory; @@ -386,12 +388,19 @@ impl Inserter { }, instant_requests: RegionInsertRequests::default(), }; + let table_info = table_infos.values().next(); + let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str()); + let schema = + table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone()); let write_cost = write_meter!( - ctx.current_catalog(), - ctx.current_schema(), + catalog, + &schema, metered, + ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?), ctx.channel() as u8 - ); + ) + .await + .context(WriteRejectedSnafu)?; prepared.retain(|(_, batch)| batch.num_rows() != 0); let results = if prepared.is_empty() { Vec::new() @@ -399,8 +408,7 @@ impl Inserter { // One original request shares admission across all table submissions. let permit = batcher.acquire().await?; let submissions = prepared.into_iter().map(|(info, batch)| { - // Routing uses the target database; metering above retains the - // original request context, including fully qualified SQL writes. + // Route to the same target database used for admission above. let mut target_ctx = ctx.fork(); target_ctx.set_current_catalog(&info.catalog_name); target_ctx.set_current_schema(&info.schema_name); @@ -542,6 +550,72 @@ impl Inserter { } } +/// Admits a finite request before it is split into internal writes. +/// The returned context preserves accounting while preventing a second row debit. +pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result { + // The zero value is WCU: this record only admits rows. Actual inserts retain + // their existing WCU accounting, so charging here would count it twice. + write_meter!(MeterRecord::new( + ctx.current_catalog().to_string(), + ctx.current_schema(), + 0, + ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows), + ctx.channel() as u8, + )) + .await + .context(WriteRejectedSnafu)?; + Ok(Arc::new(ctx.with_write_admission())) +} + +/// Admits all database totals before dispatching any batch of a finite request. +/// Each batch keeps its own protocol options and target database. +pub async fn admit_row_insert_batches( + batches: &mut [(QueryContextRef, RowInsertRequests)], +) -> Result<()> { + let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new(); + for (ctx, requests) in batches.iter() { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 { + continue; + } + let (_, total) = totals + .entry((catalog.to_string(), schema.clone())) + .or_insert_with(|| (ctx.clone(), 0)); + for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) { + *total = + total + .checked_add(rows.rows.len() as u64) + .context(InvalidInsertRequestSnafu { + reason: "Insert row count exceeds u64::MAX", + })?; + } + } + for (ctx, rows) in totals.values() { + admit_write(*rows, ctx).await?; + } + for (ctx, _) in batches { + *ctx = Arc::new(ctx.with_write_admission()); + } + Ok(()) +} + +fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result { + requests + .normal_requests + .requests + .iter() + .chain(&requests.instant_requests.requests) + .filter_map(|request| request.rows.as_ref()) + .try_fold(0u64, |total, rows| { + total + .checked_add(rows.rows.len() as u64) + .context(InvalidInsertRequestSnafu { + reason: "Insert row count exceeds u64::MAX", + }) + }) +} + impl Inserter { async fn do_request( &self, @@ -552,12 +626,21 @@ impl Inserter { // Fill impure default values in the request let requests = fill_reqs_with_impure_default(table_infos, requests)?; + // All tables in a batch resolve to the same database. Qualified SQL + // inserts may target a different database than the session's current one. + let table_info = table_infos.values().next(); + let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str()); + let schema = + table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone()); let write_cost = write_meter!( - ctx.current_catalog(), - ctx.current_schema(), + catalog, + schema.clone(), requests, + ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?), ctx.channel() as u8 - ); + ) + .await + .context(WriteRejectedSnafu)?; let request_factory = RegionRequestFactory::new(RegionRequestHeader { tracing_context: TracingContext::from_current_span().to_w3c(), dbname: ctx.get_db_string(), @@ -1721,7 +1804,9 @@ mod tests { use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType}; use crate::insert::*; - use crate::test_util::{create_partition_rule_manager, prepare_mocked_backend}; + use crate::test_util::{ + create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, + }; fn make_table_ref_with_schema( ts_name: &str, @@ -1909,6 +1994,340 @@ mod tests { assert!(table_is_native_histogram(&table)); } + // Keep global meter registration in one test, isolated by nextest's per-test process. + #[tokio::test] + async fn test_write_meter_admission() { + use std::cell::Cell; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + + use api::region::RegionResponse; + use api::v1::region::region_request::Body; + use arrow::array::{Int32Array, TimestampMillisecondArray}; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + use common_grpc::flight::{FlightEncoder, FlightMessage}; + use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; + use futures::future::BoxFuture; + use meter_core::ItemCalculator; + use meter_core::collect::{Collect, WriteRejected}; + use meter_core::data::MeterRecord; + use meter_core::global::global_registry; + use session::context::Channel; + + const CATALOG: &str = "write_meter_test"; + + #[derive(Default)] + struct Meter { + reject: AtomicBool, + attempts: Mutex>, + accepted_value: AtomicU64, + } + + impl Collect for Meter { + fn on_write( + &self, + record: MeterRecord, + ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> { + Box::pin(async move { + if record.catalog != CATALOG { + return Ok(()); + } + let value = record.value; + self.attempts.lock().unwrap().push(record); + if self.reject.load(Ordering::Relaxed) { + return Err(WriteRejected::new("database row quota exhausted")); + } + self.accepted_value.fetch_add(value, Ordering::Relaxed); + Ok(()) + }) + } + + fn on_read(&self, _: MeterRecord) {} + } + + impl ItemCalculator for Meter { + fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 { + 17 + } + } + + let kv_backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(kv_backend.clone()).await; + let (sender, mut dispatched) = tokio::sync::mpsc::channel(16); + let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| { + let rows = match request.body.unwrap() { + Body::Inserts(requests) => requests + .requests + .iter() + .filter_map(|request| request.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(), + // The bulk batches below each contain two rows. + Body::BulkInsert(_) => 2, + body => panic!("unexpected request: {body:?}"), + }; + Ok(RegionResponse::new(rows)) + }); + let flow_cache = Cache::new(100); + let inserter = Inserter::new( + catalog::memory::MemoryCatalogManager::new(), + partition_manager, + Arc::new(MockDatanodeManager::new(watcher)), + Arc::new(new_table_flownode_set_cache( + String::new(), + flow_cache.clone(), + kv_backend, + )), + true, + ); + let mut table_info = new_test_table_info(1, "table_1", [1].into_iter()); + table_info.catalog_name = CATALOG.to_string(); + table_info.schema_name = "target_db".to_string(); + let table_info = Arc::new(table_info); + let table_infos = HashMap::from_iter([(1, table_info.clone())]); + let ctx = Arc::new(QueryContext::with_channel( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Channel::Postgres, + )); + let rows_request = || { + let build = |num_rows| RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { values: vec![] }; num_rows], + }), + ..Default::default() + }], + }; + InstantAndNormalInsertRequests { + normal_requests: build(3), + instant_requests: build(2), + } + }; + + // No collector or calculator: ordinary OSS insertion still succeeds. + let output = inserter + .do_request(rows_request(), &table_infos, &ctx) + .await + .unwrap(); + assert_eq!(output.meta.cost, 0); + assert!(matches!(output.data, OutputData::AffectedRows(3))); + dispatched.try_recv().unwrap(); + + let meter = Arc::new(Meter::default()); + global_registry().set_collector(meter.clone()); + global_registry().register_calculator(meter.clone()); + // The dependency controls noop mode; exercise both builds with this test. + let enabled = Cell::new(false); + write_meter!({ + enabled.set(true); + MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0) + }) + .await + .unwrap(); + let enabled = enabled.get(); + + let output = inserter + .do_request(rows_request(), &table_infos, &ctx) + .await + .unwrap(); + assert_eq!(output.meta.cost, if enabled { 17 } else { 0 }); + assert!(matches!(output.data, OutputData::AffectedRows(3))); + dispatched.try_recv().unwrap(); + assert!(flow_cache.contains_key(&1)); + flow_cache.invalidate_all(); + meter.reject.store(true, Ordering::Relaxed); + let result = inserter + .do_request(rows_request(), &table_infos, &ctx) + .await; + if enabled { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!(error.to_string().contains("database row quota exhausted")); + assert!(dispatched.try_recv().is_err()); + assert!( + !flow_cache.contains_key(&1), + "rejected write reached flow mirroring" + ); + assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17); + let attempts = meter.attempts.lock().unwrap(); + assert_eq!(attempts.len(), 2); + for record in attempts.iter() { + assert_eq!(record.catalog, CATALOG); + assert_eq!(record.schema, "target_db"); + assert_eq!( + (record.rows, record.value, record.source), + (5, 17, Channel::Postgres as u8) + ); + } + } else { + assert_eq!(result.unwrap().meta.cost, 0); + dispatched.try_recv().unwrap(); + assert!(meter.attempts.lock().unwrap().is_empty()); + } + meter.attempts.lock().unwrap().clear(); + + // Rejection must also precede admission to the table batcher's queue. + if enabled { + let batcher: Arc = Arc::new(UnexpectedBatcher); + let rows = Rows { + schema: vec![ + api::v1::helper::tag_column_schema("a", ColumnDataType::Int32), + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + field_column_schema("b", ColumnDataType::Int32), + ], + rows: vec![api::v1::Row { + values: vec![ + api::v1::value::ValueData::I32Value(60).into(), + Value { + value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue( + 0, + )), + }, + api::v1::value::ValueData::I32Value(0).into(), + ], + }], + }; + let error = inserter + .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RateLimited); + let mut attempts = meter.attempts.lock().unwrap(); + assert_eq!(attempts.len(), 1); + assert_eq!(attempts[0].catalog, CATALOG); + assert_eq!(attempts[0].schema, "target_db"); + assert_eq!(attempts[0].rows, 1); + attempts.clear(); + } + + let table = Arc::new(table::Table::new( + table_info.clone(), + table::metadata::FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )); + let batch = RecordBatch::try_new( + table_info.meta.schema.arrow_schema().clone(), + vec![ + Arc::new(Int32Array::from(vec![60, 70])), + Arc::new(TimestampMillisecondArray::from(vec![0, 1])), + Arc::new(Int32Array::from(vec![0, 0])), + ], + ) + .unwrap(); + let bulk_insert = |batch: RecordBatch| { + let flight_data = FlightEncoder::default() + .encode(FlightMessage::RecordBatch(batch.clone())) + .into_iter() + .next() + .unwrap(); + inserter.handle_bulk_insert( + table.clone(), + flight_data, + batch, + Bytes::new(), + false, + Channel::Grpc, + ) + }; + // Empty batches bypass admission even while the collector rejects. + assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0); + assert!(meter.attempts.lock().unwrap().is_empty()); + assert!(dispatched.try_recv().is_err()); + + // A collector change between batches takes effect on the very next batch. + for reject in [false, true, false] { + meter.reject.store(reject, Ordering::Relaxed); + let result = bulk_insert(batch.clone()).await; + if enabled && reject { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!(dispatched.try_recv().is_err()); + } else { + assert_eq!(result.unwrap(), 2); + dispatched.try_recv().unwrap(); + } + } + { + let attempts = meter.attempts.lock().unwrap(); + assert_eq!(attempts.len(), if enabled { 3 } else { 0 }); + for record in attempts.iter() { + assert_eq!(record.catalog, CATALOG); + assert_eq!(record.schema, "target_db"); + assert_eq!( + (record.rows, record.value, record.source), + (2, 0, Channel::Grpc as u8) + ); + } + } + assert_eq!( + meter.accepted_value.load(Ordering::Relaxed), + if enabled { 17 } else { 0 } + ); + + // Aggregate repeated database targets and retain admission across nested + // batching without changing the caller's reusable context. + meter.attempts.lock().unwrap().clear(); + let original = Arc::new(QueryContext::with_channel( + CATALOG, + "a", + Channel::Prometheus, + )); + let mut batches = ["a", "b", "a"].map(|schema| { + let ctx = if schema == "a" { + original.clone() + } else { + Arc::new(QueryContext::with_channel( + CATALOG, + schema, + Channel::Prometheus, + )) + }; + ( + ctx, + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "data".into(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + }], + }, + ) + }); + admit_row_insert_batches(&mut batches).await.unwrap(); + admit_row_insert_batches(&mut batches).await.unwrap(); + assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4); + for (ctx, _) in &batches { + assert_eq!( + ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2), + 0 + ); + assert_eq!(ctx.channel(), Channel::Prometheus); + } + let attempts = meter.attempts.lock().unwrap(); + let totals = attempts + .iter() + .map(|r| (r.schema.as_str(), r.rows, r.value)) + .collect::>(); + assert_eq!( + totals, + if enabled { + vec![("a", 4, 0), ("b", 2, 0)] + } else { + vec![] + } + ); + } + #[test] fn test_skip_wal_does_not_change_table_options() { check_skip_wal_does_not_change_table_options(false); diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 2f020c28e68..0bebd3c1295 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -101,7 +101,7 @@ use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ AlterKind, AlterTableRequest, AnnotationContext, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, - TableOptions, validate_and_normalize_annotation_options, + INGEST_ROWS_RATE_LIMIT_KEY, TableOptions, validate_and_normalize_annotation_options, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -357,11 +357,10 @@ impl StatementExecutor { .map(|v| v.into_inner()); let create_expr = &mut expr_helper::create_to_expr(&stmt, &ctx)?; - // Don't inherit schema-level TTL/compaction options into table options: - // TTL is applied during compaction, and `compaction.*` is handled separately. + // TTL and compaction options are handled separately; ingestion quota is database-only. if let Some(schema_options) = schema_options { for (key, value) in schema_options.extra_options.iter() { - if key.starts_with("compaction.") { + if key == INGEST_ROWS_RATE_LIMIT_KEY || key.starts_with("compaction.") { continue; } create_expr diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 2a446033f62..006ac59771c 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -25,6 +25,12 @@ provided through handler traits, mostly implemented by `frontend`. after bumping it. - Auth or `QueryContext` changes must cover HTTP, gRPC, MySQL, and PostgreSQL entry points. +- Internal frontend gRPC listeners set `Channel::Internal` in server-owned request + extensions. Handlers preserve it through query execution and metering; client + flow metadata does not grant an internal channel. +- Finite requests split by wire handlers (Prometheus batching and OpenTSDB + summary/details) admit all rows before dispatch. Admission is process-local + `QueryContext` state; never accept it from headers or serialized contexts. - New protocols or externally visible routes also require frontend service and configuration wiring. diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 8ffd97e2c0f..0e2d896db8c 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -82,6 +82,8 @@ jsonb.workspace = true lazy_static.workspace = true log-query.workspace = true loki-proto.workspace = true +meter-core.workspace = true +meter-macros.workspace = true metric-engine.workspace = true mime_guess = "2.0" mysql_common = "0.38" diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index de872ac326b..9543b0d3a88 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -40,6 +40,8 @@ use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_runtime::spawn_global; use common_telemetry::{debug, error, warn}; use datatypes::timestamp::append_timestamps; +use meter_core::data::MeterRecord; +use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; @@ -276,6 +278,21 @@ impl LogicalTablePendingRowsBatcher { return Ok(0); } + // Flushes dispatch directly to datanodes, so admit once before enqueueing. + write_meter!(MeterRecord::new( + ctx.current_catalog().to_string(), + ctx.current_schema(), + 0, + ctx.write_rows_to_admit( + ctx.current_catalog(), + &ctx.current_schema(), + total_rows as u64 + ), + ctx.channel() as u8, + )) + .await + .context(error::WriteRejectedSnafu)?; + let permit = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_acquire_inflight_permit"]) diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 30997a8536a..47f9953a187 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -62,6 +62,14 @@ pub enum Error { #[snafu(display("Pending rows batcher channel closed"))] BatcherChannelClosed, + #[snafu(display("Write rejected: {error}"))] + WriteRejected { + #[snafu(source)] + error: meter_core::collect::WriteRejected, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported data type: {}, reason: {}", data_type, reason))] UnsupportedDataType { data_type: ConcreteDataType, @@ -855,7 +863,7 @@ impl ErrorExt for Error { Suspended { .. } => StatusCode::Suspended, - MemoryLimitExceeded { .. } => StatusCode::RateLimited, + MemoryLimitExceeded { .. } | WriteRejected { .. } => StatusCode::RateLimited, GreptimeProto { source, .. } => source.status_code(), Partition { source, .. } => source.status_code(), @@ -901,7 +909,7 @@ impl ErrorExt for Error { MemoryLimitExceeded { source, .. } => source.retry_hint(), CollectRecordbatch { source, .. } => source.retry_hint(), - TooManyConcurrentRequests { .. } => RetryHint::Retryable, + TooManyConcurrentRequests { .. } | WriteRejected { .. } => RetryHint::Retryable, _ => RetryHint::NonRetryable, } diff --git a/src/servers/src/grpc/context_auth.rs b/src/servers/src/grpc/context_auth.rs index 2e0bda0939c..bb6f2d27d2a 100644 --- a/src/servers/src/grpc/context_auth.rs +++ b/src/servers/src/grpc/context_auth.rs @@ -34,9 +34,10 @@ use crate::http::AUTHORIZATION_HEADER; use crate::http::header::constants::GREPTIME_DB_HEADER_NAME; use crate::metrics::METRIC_AUTH_FAILURE; -/// Create a query context from the grpc metadata. +/// Create a query context from gRPC metadata and server-owned request extensions. pub fn create_query_context_from_grpc_metadata( headers: &MetadataMap, + extensions: &http::Extensions, ) -> TonicResult { let (catalog, schema) = if let Some(db) = extract_header(headers, &[GREPTIME_DB_HEADER_NAME])? { parse_catalog_and_schema_from_db_string(db) @@ -50,7 +51,12 @@ pub fn create_query_context_from_grpc_metadata( let ctx = QueryContextBuilder::default() .current_catalog(catalog) .current_schema(schema) - .channel(Channel::Grpc) + .channel( + extensions + .get::() + .copied() + .unwrap_or(Channel::Grpc), + ) .build(); // OTEL Arrow uses ordinary inserts. Accept only its request-level WAL hint, // leaving unrelated hints and reserved internal extensions unchanged. @@ -183,11 +189,28 @@ mod tests { use super::*; + #[test] + fn test_channel_comes_from_server_extensions() { + let mut headers = MetadataMap::new(); + headers.insert( + "x-greptime-flow-extensions", + r#"[["flow.return_region_seq","true"]]"#.parse().unwrap(), + ); + headers.insert(HINTS_KEY, "channel=internal".parse().unwrap()); + let mut extensions = http::Extensions::new(); + let ctx = create_query_context_from_grpc_metadata(&headers, &extensions).unwrap(); + assert_eq!(ctx.channel(), Channel::Grpc); + + extensions.insert(Channel::Internal); + let ctx = create_query_context_from_grpc_metadata(&headers, &extensions).unwrap(); + assert_eq!(ctx.channel(), Channel::Internal); + } + #[test] fn test_arrow_insert_hint_does_not_accept_reserved_extensions() { let mut headers = MetadataMap::new(); assert_eq!( - create_query_context_from_grpc_metadata(&headers) + create_query_context_from_grpc_metadata(&headers, &Default::default()) .unwrap() .extension(INSERT_SKIP_WAL_HINT), None @@ -198,7 +221,8 @@ mod tests { hints.push_str(&format!(",{key}=external")); } headers.insert(HINTS_KEY, hints.parse().unwrap()); - let ctx = create_query_context_from_grpc_metadata(&headers).unwrap(); + let ctx = + create_query_context_from_grpc_metadata(&headers, &Default::default()).unwrap(); assert_eq!(ctx.skip_wal(), expected); assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); assert_eq!(ctx.extension("ttl"), None); @@ -219,7 +243,8 @@ mod tests { ("insert_skip_wal=true,insert_skip_wal=invalid", true), ] { headers.insert(HINTS_KEY, hints.parse().unwrap()); - let ctx = create_query_context_from_grpc_metadata(&headers).unwrap(); + let ctx = + create_query_context_from_grpc_metadata(&headers, &Default::default()).unwrap(); assert_eq!(ctx.skip_wal(), expected); } for value in ["", "TRUE", "1", "invalid"] { @@ -227,7 +252,9 @@ mod tests { HINTS_KEY, format!("insert_skip_wal={value}").parse().unwrap(), ); - assert!(create_query_context_from_grpc_metadata(&headers).is_err()); + assert!( + create_query_context_from_grpc_metadata(&headers, &Default::default()).is_err() + ); } } } diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs index 79904806aba..2a63a982bef 100644 --- a/src/servers/src/grpc/database.rs +++ b/src/servers/src/grpc/database.rs @@ -21,6 +21,7 @@ use common_query::OutputData; use common_telemetry::{debug, warn}; use futures::StreamExt; use prost::Message; +use session::context::Channel; use tonic::{Request, Response, Status, Streaming}; use crate::grpc::greptime_handler::GreptimeRequestHandler; @@ -46,6 +47,11 @@ impl GreptimeDatabase for DatabaseService { ) -> TonicResult> { let remote_addr = request.remote_addr(); let hints = hint_headers::extract_hints(request.metadata()); + let channel = request + .extensions() + .get::() + .copied() + .unwrap_or(Channel::Grpc); debug!( "GreptimeDatabase::Handle: request from {:?} with hints: {:?}", remote_addr, hints @@ -61,7 +67,7 @@ impl GreptimeDatabase for DatabaseService { let handler = self.handler.clone(); let request_future = async move { let request = request.into_inner(); - let output = handler.handle_request(request, hints).await?; + let output = handler.handle_request(request, hints, channel).await?; let message = match output.data { OutputData::AffectedRows(rows) => GreptimeResponse { header: Some(ResponseHeader { @@ -100,6 +106,11 @@ impl GreptimeDatabase for DatabaseService { ) -> Result, Status> { let remote_addr = request.remote_addr(); let hints = hint_headers::extract_hints(request.metadata()); + let channel = request + .extensions() + .get::() + .copied() + .unwrap_or(Channel::Grpc); debug!( "GreptimeDatabase::HandleRequests: request from {:?} with hints: {:?}", remote_addr, hints @@ -121,7 +132,9 @@ impl GreptimeDatabase for DatabaseService { } else { None }; - let output = handler.handle_request(request, hints.clone()).await?; + let output = handler + .handle_request(request, hints.clone(), channel) + .await?; match output.data { OutputData::AffectedRows(rows) => affected_rows += rows, OutputData::Stream(_) | OutputData::RecordBatches(_) => { diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index f0035ed3f8c..874e98b7326 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -199,12 +199,17 @@ impl FlightCraft for GreptimeRequestHandler { let mut hints = hint_headers::extract_hints(request.metadata()); hints.extend(extract_flow_extensions(request.metadata())?); let snapshot_seqs = extract_snapshot_seqs(request.metadata())?; + let channel = request + .extensions() + .get::() + .copied() + .unwrap_or(Channel::Grpc); let ticket = request.into_inner().ticket; let request = GreptimeRequest::decode(ticket.as_ref()).context(error::InvalidFlightTicketSnafu)?; let query_ctx = - create_query_context(Channel::Grpc, request.header.as_ref(), hints, snapshot_seqs)?; + create_query_context(channel, request.header.as_ref(), hints, snapshot_seqs)?; // Validate flow hint syntax at the transport boundary before dispatching the request. // This does not authorize or execute anything; `handle_request()` below still performs // the normal frontend handling and auth checks before query execution. @@ -249,7 +254,8 @@ impl FlightCraft for GreptimeRequestHandler { let limiter = extensions.get::().cloned(); - let query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers)?; + let query_ctx = + context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?; context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?; const MAX_PENDING_RESPONSES: usize = 32; diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index dcd8c641a64..ad0a693ac59 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -78,9 +78,10 @@ impl GreptimeRequestHandler { &self, request: GreptimeRequest, hints: Vec<(String, String)>, + channel: Channel, ) -> Result { let header = request.header.as_ref(); - let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?; + let query_ctx = create_query_context(channel, header, hints, HashMap::new())?; let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", })?; diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 92bd5d722e5..abec69604f6 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -47,6 +47,11 @@ pub struct PrometheusGatewayService { impl PrometheusGateway for PrometheusGatewayService { async fn handle(&self, req: Request) -> TonicResult> { let mut is_range_query = false; + let channel = req + .extensions() + .get::() + .copied() + .unwrap_or(Channel::Promql); let inner = req.into_inner(); let prom_query = match inner.promql.context(InvalidQuerySnafu { reason: "Expecting non-empty PromqlRequest.", @@ -80,12 +85,8 @@ impl PrometheusGateway for PrometheusGatewayService { }; let header = inner.header.as_ref(); - let query_ctx = create_query_context( - Channel::Promql, - header, - Default::default(), - Default::default(), - )?; + let query_ctx = + create_query_context(channel, header, Default::default(), Default::default())?; let user_info = auth(self.user_provider.clone(), header, &query_ctx).await?; query_ctx.set_current_user(user_info); diff --git a/src/servers/src/http/opentsdb.rs b/src/servers/src/http/opentsdb.rs index 51667db2fe7..12937a4d184 100644 --- a/src/servers/src/http/opentsdb.rs +++ b/src/servers/src/http/opentsdb.rs @@ -88,22 +88,22 @@ pub async fn put( .collect::>(); ctx.set_channel(Channel::Opentsdb); - let ctx = Arc::new(ctx); + let mut ctx = Arc::new(ctx); if summary || details { opentsdb_handler .preflight(&data_points, ctx.clone()) .await?; + ctx = operator::insert::admit_write(data_points.len() as u64, &ctx) + .await + .map_err(common_error::ext::BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; } let response = if !summary && !details { - if let Err(e) = opentsdb_handler.exec_batch(data_points, ctx.clone()).await { - // Not debugging purpose, failed fast. - return error::InternalSnafu { - err_msg: e.to_string(), - } - .fail(); - } + opentsdb_handler + .exec_batch(data_points, ctx.clone()) + .await?; (HttpStatusCode::NO_CONTENT, Json(OpentsdbPutResponse::Empty)) } else { let mut response = OpentsdbDebuggingResponse { diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index ed1c778f701..17db6abae04 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -380,12 +380,16 @@ async fn preflight_prometheus_rows( prom_store_handler: &PromStoreProtocolHandlerRef, batches: &mut [PromWriteBatch], ) -> Result<()> { - for (ctx, reqs) in batches { + for (ctx, reqs) in batches.iter_mut() { prom_store_handler.pre_write(reqs, ctx.clone()).await?; // Detach from context clones retained by pre-write hooks so the checked // schema cannot change before this prepared batch is written. *ctx = Arc::new(ctx.fork()); } + operator::insert::admit_row_insert_batches(batches) + .await + .map_err(common_error::ext::BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; Ok(()) } @@ -918,7 +922,8 @@ mod tests { #[async_trait] impl PromWriteBatcher for RecordingPromWriteBatcher { - async fn submit(&self, requests: RowInsertRequests, _ctx: QueryContextRef) -> Result { + async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { + assert_eq!(ctx.write_rows_to_admit("greptime", "public", 1), 0); record_write_event(&self.events, "batch", &requests); Ok(prom_write_row_count(&requests)) } @@ -942,9 +947,10 @@ mod tests { async fn write_prepared( &self, request: RowInsertRequests, - _ctx: QueryContextRef, + ctx: QueryContextRef, _with_metric_engine: bool, ) -> Result { + assert_eq!(ctx.write_rows_to_admit("greptime", "public", 1), 0); record_write_event(&self.events, "direct", &request); Ok(Output::new_with_affected_rows(0)) } diff --git a/src/servers/src/otel_arrow.rs b/src/servers/src/otel_arrow.rs index 730020161d3..0cf07c76468 100644 --- a/src/servers/src/otel_arrow.rs +++ b/src/servers/src/otel_arrow.rs @@ -84,9 +84,14 @@ impl ArrowMetricsService for OtelArrowServiceHandler Result, Status> { let (mut sender, receiver) = futures::channel::mpsc::channel(100); - let (headers, _, mut incoming_requests) = request.into_parts(); + let (headers, extensions, mut incoming_requests) = request.into_parts(); - let query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers)?; + // OTEL Arrow is currently used only for external ingestion. Its service bypasses + // the frontend router middleware, so even the internal listener defaults to + // Channel::Grpc and remains rate-limited. Before using this path internally, + // propagate the server-owned Channel::Internal marker to this service. + let query_ctx = + context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?; context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?; let query_ctx = { let mut ctx = query_ctx.fork(); diff --git a/src/servers/tests/http/opentsdb_test.rs b/src/servers/tests/http/opentsdb_test.rs index eb2f312f94a..6dc2ec51d52 100644 --- a/src/servers/tests/http/opentsdb_test.rs +++ b/src/servers/tests/http/opentsdb_test.rs @@ -28,6 +28,7 @@ use servers::opentsdb::codec::DataPoint; use servers::query_handler::OpentsdbProtocolHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContextRef; +use snafu::IntoError; use sql::statements::statement::Statement; use tokio::sync::mpsc; @@ -43,6 +44,11 @@ impl OpentsdbProtocolHandler for DummyInstance { async fn exec(&self, data_points: Vec, _ctx: QueryContextRef) -> Result { let data_point = data_points.first().unwrap(); + if data_point.metric() == "rate_limited" { + return Err(error::WriteRejectedSnafu.into_error( + meter_core::collect::WriteRejected::new("database row quota exhausted"), + )); + } if data_point.metric() == "should_failed" { return error::InternalSnafu { err_msg: "expected", @@ -156,6 +162,14 @@ async fn test_opentsdb_put() { assert_eq!(result.status(), 500); assert_eq!(result.text().await, "{\"error\":\"Internal error: 1003\"}"); + let result = client + .post("/v1/opentsdb/api/put") + .body(create_data_point("rate_limited")) + .send() + .await; + assert_eq!(result.status(), 429); + assert!(result.text().await.contains("database row quota exhausted")); + let mut metrics = vec![]; while let Ok(s) = rx.try_recv() { metrics.push(s); diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 2833abd2578..0bb3f0f6a31 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -85,6 +85,9 @@ pub struct QueryContext { /// Track which protocol the query comes from. #[builder(default)] channel: Channel, + /// Process-local admission for one database in this request. Never sent over the wire. + #[builder(setter(skip))] + admitted_write: Option<(String, String)>, /// Process id for managing on-going queries #[builder(default)] process_id: u32, @@ -263,6 +266,28 @@ impl QueryContext { fork } + /// Forks a context for internal writes after admitting the complete request. + /// Call only after write admission succeeds for the current database. + pub fn with_write_admission(&self) -> Self { + let mut ctx = self.fork(); + ctx.admitted_write = Some((ctx.current_catalog().to_string(), ctx.current_schema())); + ctx + } + + /// Returns zero for writes covered by this request's database admission. + /// Usage accounting and the original protocol channel remain unchanged. + pub fn write_rows_to_admit(&self, catalog: &str, schema: &str, rows: u64) -> u64 { + if self + .admitted_write + .as_ref() + .is_some_and(|(c, s)| c == catalog && s == schema) + { + 0 + } else { + rows + } + } + pub fn arc() -> QueryContextRef { Arc::new( QueryContextBuilder::default() @@ -608,6 +633,7 @@ impl QueryContextBuilder { .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())), channel, batching_enabled: self.batching_enabled.unwrap_or_default(), + admitted_write: None, process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), protocol_ctx: self.protocol_ctx.unwrap_or_default(), @@ -788,6 +814,26 @@ mod test { assert!(fork.skip_wal()); } + #[test] + fn test_write_admission_is_local_and_database_scoped() { + let ctx = QueryContext::with_channel("greptime", "public", Channel::Otlp); + let admitted = ctx.with_write_admission(); + assert_eq!(ctx.write_rows_to_admit("greptime", "public", 100), 100); + assert_eq!(admitted.write_rows_to_admit("greptime", "public", 100), 0); + assert_eq!(admitted.write_rows_to_admit("greptime", "other", 100), 100); + assert_eq!(admitted.write_rows_to_admit("other", "public", 100), 100); + assert_eq!(admitted.channel(), Channel::Otlp); + assert_eq!( + admitted + .fork() + .write_rows_to_admit("greptime", "public", 100), + 0 + ); + let wire: api::v1::QueryContext = admitted.into(); + let restored = QueryContext::from(wire); + assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100); + } + #[test] fn test_skip_wal_is_not_serialized_in_query_context() { let context = QueryContextBuilder::default().skip_wal(true).build(); @@ -808,7 +854,7 @@ mod test { current_schema: "s1".to_string(), timezone: "UTC".to_string(), extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]), - channel: Channel::Grpc as u32, + channel: Channel::Internal as u32, snapshot_seqs: Some(api::v1::SnapshotSequences { snapshot_seqs: HashMap::from([(1, 100)]), sst_min_sequences: HashMap::from([(1, 90)]), diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 82dcd949ae4..bb73870b45f 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -1591,6 +1591,18 @@ mod tests { } _ => unreachable!(), } + + let sql = "CREATE DATABASE prometheus with ('ingest_rows_rate_limit'='1000');"; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); + let stmts = result.unwrap(); + match &stmts[0] { + Statement::CreateDatabase(c) => { + assert_eq!(c.name.to_string(), "prometheus"); + assert_eq!(c.options.get("ingest_rows_rate_limit").unwrap(), "1000"); + } + _ => unreachable!(), + } } #[test] diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index c57dc0b009d..10055df4b45 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -106,6 +106,9 @@ pub const DDL_WAIT: &str = "wait"; pub const VALID_DDL_OPTION_KEYS: [&str; 2] = [DDL_TIMEOUT, DDL_WAIT]; +/// The key of ingest rows rate limit option (rows per second, cluster-wide) in database options. +pub const INGEST_ROWS_RATE_LIMIT_KEY: &str = "ingest_rows_rate_limit"; + // Valid option keys when creating a db. static VALID_DB_OPT_KEYS: Lazy> = Lazy::new(|| { let mut set = HashSet::new(); @@ -129,6 +132,7 @@ static VALID_DB_OPT_KEYS: Lazy> = Lazy::new(|| { set.insert(TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER); set.insert(TWCS_MAX_OUTPUT_FILE_SIZE); set.insert(SST_FORMAT_KEY); + set.insert(INGEST_ROWS_RATE_LIMIT_KEY); set }); @@ -142,6 +146,12 @@ pub fn validate_database_option_value( key: &str, value: Option<&str>, ) -> std::result::Result<(), &'static str> { + if key == INGEST_ROWS_RATE_LIMIT_KEY { + return value + .and_then(|value| value.parse::().ok()) + .map(|_| ()) + .ok_or("expected a non-negative integer fitting in u64"); + } let (minimum, constraint) = match key { TWCS_TRIGGER_FILE_NUM | TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM @@ -933,6 +943,8 @@ mod tests { assert!(validate_database_option( "compaction.twcs.inactive_window.l1_merge_trigger" )); + assert!(validate_database_option(INGEST_ROWS_RATE_LIMIT_KEY)); + assert!(validate_database_option("ingest_rows_rate_limit")); assert!(!validate_database_option("foo")); } @@ -994,6 +1006,27 @@ mod tests { } } + #[test] + fn test_database_ingest_rate_limit_value_boundaries() { + for invalid in [ + None, + Some(""), + Some("abc"), + Some("1000/s"), + Some("-1"), + Some("1.5"), + Some("18446744073709551616"), + ] { + assert!(validate_database_option_value(INGEST_ROWS_RATE_LIMIT_KEY, invalid).is_err()); + } + let maximum = u64::MAX.to_string(); + for valid in ["0", "1", maximum.as_str()] { + assert!( + validate_database_option_value(INGEST_ROWS_RATE_LIMIT_KEY, Some(valid)).is_ok() + ); + } + } + #[test] fn test_serialize_table_options() { let options = TableOptions { diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index f5e94319bac..dc4d3898860 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -325,6 +325,57 @@ PARTITION ON COLUMNS (n) ( check_output_stream(output, expected).await; } +#[apply(both_instances_cases)] +async fn test_database_ingest_rate_limit_not_inherited(instance: Arc) { + let frontend = instance.frontend(); + execute_sql( + &frontend, + "CREATE DATABASE limited WITH ('ingest_rows_rate_limit'='1000', 'skip_wal'='true')", + ) + .await; + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, "limited")); + + for (name, sql) in [ + ("source", "CREATE TABLE source (ts TIMESTAMP TIME INDEX)"), + ("copy", "CREATE TABLE copy LIKE source"), + ] { + execute_sql_with(&frontend, sql, ctx.clone()).await; + let table = frontend + .catalog_manager() + .table(DEFAULT_CATALOG_NAME, "limited", name, None) + .await + .unwrap() + .unwrap(); + let options = &table.table_info().meta.options; + assert!(!options.extra_options.contains_key("ingest_rows_rate_limit")); + assert!(options.skip_wal); + + let output = + execute_sql_with(&frontend, &format!("SHOW CREATE TABLE {name}"), ctx.clone()).await; + let OutputData::RecordBatches(batches) = output.data else { + unreachable!() + }; + let batch = batches.iter().next().unwrap(); + let ddl = batch + .column_by_name("Create Table") + .unwrap() + .as_string::() + .value(0); + assert!(!ddl.contains("ingest_rows_rate_limit")); + execute_sql_with(&frontend, &format!("DROP TABLE {name}"), ctx.clone()).await; + execute_sql_with(&frontend, ddl, ctx.clone()).await; + } + + let output = execute_sql(&frontend, "SHOW CREATE DATABASE limited").await; + assert!( + output + .data + .pretty_print() + .await + .contains("ingest_rows_rate_limit") + ); +} + #[apply(standalone_instance_case)] async fn test_extra_external_table_options(instance: Arc) { let frontend = instance.frontend(); From 6a084d9a4f50d9b4f671ce35a86ceaaf9a8b90e9 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:47:23 +0000 Subject: [PATCH 39/63] fix: release completed SST write buffers during compaction (#9243) * fix: upgrade OpenDAL and honor SST write buffer size Signed-off-by: Lei, HUANG * docs: track removal of the OpenDAL bridge fork Signed-off-by: Lei, HUANG * fix(mito2): honor SST write buffers across upload paths Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- Cargo.lock | 111 +++++++++++++++------------- Cargo.toml | 6 +- src/mito2/src/access_layer.rs | 108 +++++++++++++++++++++++---- src/mito2/src/cache/write_cache.rs | 94 ++++++++++++++++++----- src/mito2/src/config.rs | 1 + src/mito2/src/flush.rs | 10 ++- src/mito2/src/sst/index.rs | 7 +- src/mito2/src/sst/parquet.rs | 27 +++++-- src/mito2/src/sst/parquet/writer.rs | 17 +++-- src/mito2/src/test_util/sst_util.rs | 71 +++++++++++++++++- src/object-store/Cargo.toml | 4 +- src/object-store/src/layers/mock.rs | 13 +++- src/object-store/src/secure_fs.rs | 45 +++++------ 13 files changed, 381 insertions(+), 133 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index de3bc7ca765..36ee6d77ffd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -912,6 +912,12 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "asyncband" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2d85fd3d291fabcc40c7232c92c280ec1754fd7b5d7ea769f143222143e179a" + [[package]] name = "asynchronous-codec" version = "0.7.0" @@ -2266,7 +2272,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -7933,7 +7939,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -9610,8 +9616,7 @@ dependencies = [ [[package]] name = "object_store_opendal" version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88f165780495c17aa3ce86846600504198c3fffd99073521552751c2430fa6ac" +source = "git+https://github.com/GreptimeTeam/opendal.git?rev=c03e56802b47de3ccfc8fe028e59a459b278a0d2#c03e56802b47de3ccfc8fe028e59a459b278a0d2" dependencies = [ "async-trait", "bytes", @@ -9698,9 +9703,9 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "opendal" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f20562cc7447fcc915fc5c23df305a412ea80a733c9f2fd9e2d267e2815be6d" +checksum = "9fe43e16d96bed57937eb7c4a28559bf4c5fc4a3951656505d1fdd4f856349b4" dependencies = [ "ctor", "opendal-core", @@ -9722,11 +9727,12 @@ dependencies = [ [[package]] name = "opendal-core" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec75551ff4cf3e57da98979f6a937aaa9ddb3915bf68cc17d03df733be6646ed" +checksum = "de4566a412776e3f65d53dd9fceced5b432a0a9c9a1e8d50ccf58823795e5c5b" dependencies = [ "anyhow", + "asyncband", "base64 0.23.1", "bytes", "futures", @@ -9734,7 +9740,6 @@ dependencies = [ "jiff", "log", "md-5 0.11.0", - "mea", "percent-encoding", "quick-xml 0.41.0", "reqsign-core", @@ -9748,9 +9753,9 @@ dependencies = [ [[package]] name = "opendal-http-transport-reqwest" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad4d4f19c3ce01126a30611f8e544eaa217104a278c889ac17c9374fe4f9e4ef" +checksum = "f772ce6c137ab43647726116d505a649a88d41928fcd813438eb30b02312bcbb" dependencies = [ "bytes", "futures", @@ -9762,21 +9767,21 @@ dependencies = [ [[package]] name = "opendal-layer-concurrent-limit" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "249ac5b0aa5a7a6c3737342d10456067937f9c9a6f3f02544271f7908ab91081" +checksum = "3734c648617e5ad724d49175dde8323550624f9614f50e9e81b7aafae30b1734" dependencies = [ + "asyncband", "futures", "http 1.5.0", - "mea", "opendal-core", ] [[package]] name = "opendal-layer-logging" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c75411ab00f77851ff086b686c1e9ca8175ac18c15afa2cb75b9036436cb06c" +checksum = "0a8daab5a91eac76f94516e7c71e677195c9221840ce8f5e71815e3fc8207742" dependencies = [ "log", "opendal-core", @@ -9784,9 +9789,9 @@ dependencies = [ [[package]] name = "opendal-layer-observe-metrics-common" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12dbc31b27ee8e5f658ae6c8514b8276059e5235e2f85138a42e6ed172324876" +checksum = "0fbe664be117c293cfa3980a0d2fc238bbb38eae5f6c191f376e88ccb830d26b" dependencies = [ "futures", "http 1.5.0", @@ -9795,9 +9800,9 @@ dependencies = [ [[package]] name = "opendal-layer-prometheus" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c14e6cb0b3ea8d83bef751613b8bf4512678afe8b49c2c96a633515bc0b34d79" +checksum = "b32078d8248555ee6271332ec480b78306a44f0d902ea4721ccac00d9b866a21" dependencies = [ "opendal-core", "opendal-layer-observe-metrics-common", @@ -9806,9 +9811,9 @@ dependencies = [ [[package]] name = "opendal-layer-retry" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80b7738bd5f233ad8da39af9b9316b9b7a4eaddd91e8e32a1e19b7030688121d" +checksum = "eb8e22cb1f86ef9af771a8947d3ccf380c8f5b76236f628314dd41c9ba73331a" dependencies = [ "backon", "log", @@ -9817,9 +9822,9 @@ dependencies = [ [[package]] name = "opendal-layer-timeout" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a704141924500f3803c05ed871b53305d2a2f11cb5ef20160c3ee688a1857f66" +checksum = "08ac6c77fe0d1b18e0064a02c19b98dbebcab78bf5e7753ad1f3abbc18850ff4" dependencies = [ "opendal-core", "tokio", @@ -9827,9 +9832,9 @@ dependencies = [ [[package]] name = "opendal-layer-tracing" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4827cf8b2ebd9123b1fb7286c421be079948e2b963a452f1f7e8aee2c12e7997" +checksum = "603807927a799ea981f62a3b0012e17181e39f468793b78ca8bedda66d19813c" dependencies = [ "futures", "http 1.5.0", @@ -9839,9 +9844,9 @@ dependencies = [ [[package]] name = "opendal-service-azblob" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3310fbbb48f111c6f590473c2cd15e1b7f8e384444b0d4e328f0464c864d767" +checksum = "c6cf31b02e7b44191d97c07254eea109b5da9c17210a29e478c8d0dc26dffb1f" dependencies = [ "base64 0.23.1", "bytes", @@ -9860,9 +9865,9 @@ dependencies = [ [[package]] name = "opendal-service-azure-common" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7348c88edf15af435b7be930077746b569fac5e738c1bf6a363b675e7317c9df" +checksum = "8107324c3f3cb9970fe20526c2c341b7c5769acbd4699bab73bd525f5fd63d62" dependencies = [ "http 1.5.0", "opendal-core", @@ -9870,9 +9875,9 @@ dependencies = [ [[package]] name = "opendal-service-fs" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "826c4e17a30643b888fe983897f9a4b23b07066e1d069727a923cc8fb419a702" +checksum = "a0e7d8974aebd33be899200c9d0d50b2141fd5fa5c462b29e20653e3ddbfe339" dependencies = [ "bytes", "log", @@ -9884,9 +9889,9 @@ dependencies = [ [[package]] name = "opendal-service-gcs" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "007f3fba63c21e516c956b891e96ff9892d8175662bfb781cdada9d3766a11e6" +checksum = "3d80eca6f9227bdf65d43e5b8d5fe75dca8cfab1da07571442a0420400847bea" dependencies = [ "async-trait", "bytes", @@ -9901,13 +9906,14 @@ dependencies = [ "serde", "serde_json", "tokio", + "uuid", ] [[package]] name = "opendal-service-http" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2fd1f0f28bd052e068d9bb4bc5d30a09de22a4f7cea24558dbeef2b91090cc3" +checksum = "9ac1553d6b57ca79e200e5c7ded9ab708f2016dadb28ede0f2da71e02fd0b960" dependencies = [ "http 1.5.0", "log", @@ -9917,11 +9923,11 @@ dependencies = [ [[package]] name = "opendal-service-mysql" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d4a25f582068ce30ba24eb3efd4ac16ca833c87b5282d8d79634fdbead74d84" +checksum = "e15592be32b105cd152ab398b3db847e22c8c5f230d9b34ff095ea3d0dc6fd97" dependencies = [ - "mea", + "asyncband", "opendal-core", "serde", "sqlx", @@ -9929,9 +9935,9 @@ dependencies = [ [[package]] name = "opendal-service-oss" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd528ec2d49c5ca69e674ffed7b3e0686fb9cfcfea0596870de381467fda4f1b" +checksum = "a0ebfa355b1d4edb0e5b2a5b86d8de9f811fae1d7baca6d65246b230b2afc0bd" dependencies = [ "bytes", "http 1.5.0", @@ -9946,9 +9952,9 @@ dependencies = [ [[package]] name = "opendal-service-s3" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58e80cdf192d7eff05feed747894d64f81905ac4eaf132edf7ea270abdd2d663" +checksum = "2942f2d8d3d4953c0e8d07cd1244879628d084948e48ce6b6d808c5938f7858d" dependencies = [ "base64 0.23.1", "bytes", @@ -12420,9 +12426,9 @@ dependencies = [ [[package]] name = "reqsign-aws-core" -version = "3.0.3" +version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4af084e1f3cbf3e67e0c972765399bce54ecec804cceba46b39a8331f3c1bff" +checksum = "bac4749b7dfa7bfaccd01eb03e9dc795ed37e3f20d6f0f38e2c67ee85ad6bc86" dependencies = [ "bytes", "form_urlencoded", @@ -12441,9 +12447,9 @@ dependencies = [ [[package]] name = "reqsign-aws-v4" -version = "3.1.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ac5b3b7cefa28933792b439186459f77f19f9b6edbeab41b8b187150361a206" +checksum = "ff250f0fd0b913fbd565e405acc553da0f13bde30bfb5403178c9d0313cdc15f" dependencies = [ "bytes", "http 1.5.0", @@ -12477,9 +12483,9 @@ dependencies = [ [[package]] name = "reqsign-core" -version = "3.2.1" +version = "3.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c07dd510b1e1b9b241883e483358147fb2ed2d497a7b39b065ba61eb93deceb0" +checksum = "ff052daffb0599681c50f85c59e7236438976efe991ab864edd9f3b235501a0f" dependencies = [ "anyhow", "base64 0.23.1", @@ -12490,6 +12496,7 @@ dependencies = [ "http 1.5.0", "jiff", "log", + "mea", "percent-encoding", "rsa", "serde", @@ -12501,9 +12508,9 @@ dependencies = [ [[package]] name = "reqsign-file-read-tokio" -version = "3.0.4" +version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "663d9d55abd0df0830ef0ae43708297cc1371cf4e8ca91f3ac813c309cca8c98" +checksum = "b3235df90a6bca681aa47dd86f2393d122a6d77042aa8a7c81e218cd45c5bfc0" dependencies = [ "anyhow", "reqsign-core", @@ -16796,7 +16803,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fa01426bbbf..b4c4c9d7f54 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -185,7 +185,11 @@ nalgebra = "0.33" nix = { version = "0.30.1", default-features = false, features = ["event", "fs", "process"] } notify = "8.0" num_cpus = "1.16" -object_store_opendal = "0.58" +# Keep object_store 0.13 compatibility while using OpenDAL 0.59. +# TODO(#9240): Switch back to the official bridge after the DataFusion/Arrow/Parquet +# upgrade aligns object_store versions. Keep the bridge and OpenDAL releases compatible. +# https://github.com/GreptimeTeam/greptimedb/issues/9240 +object_store_opendal = { git = "https://github.com/GreptimeTeam/opendal.git", rev = "c03e56802b47de3ccfc8fe028e59a459b278a0d2" } once_cell = "1.18" opentelemetry-proto = { version = "0.32", features = [ "gen-tonic", diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index ac18f838db2..337386aa2ff 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use async_stream::try_stream; +use common_base::readable_size::ReadableSize; use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; @@ -49,7 +50,7 @@ use crate::sst::location::{self, region_dir_from_table_dir}; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::sst::parquet::writer::ParquetWriter; use crate::sst::parquet::{SstInfo, WriteOptions}; -use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FormatType}; +use crate::sst::{DEFAULT_WRITE_CONCURRENCY, FormatType}; pub type AccessLayerRef = Arc; /// SST write results. @@ -244,6 +245,7 @@ impl AccessLayer { } let attempted_files = file_ids.to_vec(); + // Deleter does not normalize leading slashes like Operator::delete does. let paths: Vec<_> = file_ids .iter() .map(|file_id| { @@ -252,6 +254,8 @@ impl AccessLayer { RegionFileId::new(region_id, *file_id), self.path_type, ) + .trim_start_matches('/') + .to_string() }) .collect(); @@ -264,7 +268,7 @@ impl AccessLayer { file_ids: attempted_files.clone(), })?; deleter - .delete_iter(paths.iter().map(String::as_str)) + .delete_iter(paths) .await .with_context(|_| DeleteSstsSnafu { region_id, @@ -292,7 +296,11 @@ impl AccessLayer { .collect(); let paths: Vec<_> = index_ids .iter() - .map(|index_id| location::index_file_path(&self.table_dir, *index_id, self.path_type)) + .map(|index_id| { + location::index_file_path(&self.table_dir, *index_id, self.path_type) + .trim_start_matches('/') + .to_string() + }) .collect(); let mut deleter = self @@ -303,7 +311,7 @@ impl AccessLayer { file_ids: file_ids.clone(), })?; deleter - .delete_iter(paths.iter().map(String::as_str)) + .delete_iter(paths) .await .context(DeleteIndexesSnafu { file_ids: file_ids.clone(), @@ -468,9 +476,9 @@ impl AccessLayer { pub(crate) async fn put_sst( &self, data: &bytes::Bytes, - region_id: RegionId, - sst_info: &SstInfo, + region_file_id: RegionFileId, cache_manager: &CacheManagerRef, + write_buffer_size: ReadableSize, ) -> Result { if let Some(write_cache) = cache_manager.write_cache() { // Write to cache and upload to remote store @@ -482,27 +490,27 @@ impl AccessLayer { remote_store: self.object_store.clone(), }; write_cache - .put_and_upload_sst(data, region_id, sst_info, upload_request) + .put_and_upload_sst(data, region_file_id, upload_request, write_buffer_size) .await } else { let start = Instant::now(); - let cleaner = TempFileCleaner::new(region_id, self.object_store.clone()); + let cleaner = + TempFileCleaner::new(region_file_id.region_id(), self.object_store.clone()); let path_provider = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type); - let sst_file_path = - path_provider.build_sst_file_path(RegionFileId::new(region_id, sst_info.file_id)); + let sst_file_path = path_provider.build_sst_file_path(region_file_id); let mut writer = self .object_store .writer_with(&sst_file_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(write_buffer_size.as_bytes() as usize) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .context(OpenDalSnafu)?; if let Err(err) = writer.write(data.clone()).await.context(OpenDalSnafu) { - cleaner.clean_by_file_id(sst_info.file_id).await; + cleaner.clean_by_file_id(region_file_id.file_id()).await; return Err(err); } if let Err(err) = writer.close().await.context(OpenDalSnafu) { - cleaner.clean_by_file_id(sst_info.file_id).await; + cleaner.clean_by_file_id(region_file_id.file_id()).await; return Err(err); } let mut metrics = Metrics::new(WriteType::Flush); @@ -770,7 +778,81 @@ impl FilePathProvider for RegionFilePathFactory { mod tests { use std::path::PathBuf; + use bytes::Bytes; + use common_test_util::temp_dir::create_temp_dir; + use super::*; + use crate::cache::CacheManager; + use crate::cache::test_util::new_fs_store; + use crate::test_util::TestEnv; + use crate::test_util::sst_util::WriteChunkRecorder; + + #[rstest::rstest] + #[tokio::test] + async fn test_put_sst_uses_configured_buffer_size( + #[values(false, true)] enable_write_cache: bool, + #[values(1024, 4096)] chunk_size: usize, + ) { + let mut env = TestEnv::new().await; + let remote_chunks = WriteChunkRecorder::default(); + let remote_store = env.init_object_store_manager().layer(remote_chunks.layer()); + let local_dir = create_temp_dir("encoded-sst-cache"); + let local_chunks = WriteChunkRecorder::default(); + let local_store = + new_fs_store(local_dir.path().to_str().unwrap()).layer(local_chunks.layer()); + let write_cache = if enable_write_cache { + Some( + env.create_write_cache(local_store.clone(), ReadableSize::mb(10)) + .await, + ) + } else { + None + }; + let cache_manager = Arc::new( + CacheManager::builder() + .write_cache(write_cache.clone()) + .build(), + ); + let access_layer = AccessLayer::new( + "test", + PathType::Bare, + remote_store.clone(), + env.get_puffin_manager(), + env.get_intermediate_manager(), + ); + let file_id = RegionFileId::new(RegionId::new(1024, 1), FileId::random()); + // Two full chunks and a partial tail distinguish the configured size from the default. + let encoded = Bytes::from( + (0..2 * chunk_size + 17) + .map(|i| (i % 251) as u8) + .collect::>(), + ); + access_layer + .put_sst( + &encoded, + file_id, + &cache_manager, + ReadableSize(chunk_size as u64), + ) + .await + .unwrap(); + + let remote_path = location::sst_file_path("test", file_id, PathType::Bare); + remote_chunks.assert_chunks(&remote_path, chunk_size, encoded.len()); + assert_eq!( + remote_store.read(&remote_path).await.unwrap().to_bytes(), + encoded + ); + if let Some(write_cache) = write_cache { + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + let cache_path = write_cache.file_cache().cache_file_path(key); + local_chunks.assert_chunks(&cache_path, chunk_size, encoded.len()); + assert_eq!( + local_store.read(&cache_path).await.unwrap().to_bytes(), + encoded + ); + } + } #[tokio::test] async fn test_new_fs_cache_store() { diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 03c7bdddd73..0e96c2f16a2 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -38,8 +38,8 @@ use crate::sst::file::RegionFileId; use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; +use crate::sst::parquet::WriteOptions; use crate::sst::parquet::writer::ParquetWriter; -use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; /// Wraps the remote object store used by write cache uploads. @@ -176,11 +176,12 @@ impl WriteCache { pub(crate) async fn put_and_upload_sst( &self, data: &bytes::Bytes, - region_id: RegionId, - sst_info: &SstInfo, + region_file_id: RegionFileId, upload_request: SstUploadRequest, + write_buffer_size: ReadableSize, ) -> Result { - let file_id = sst_info.file_id; + let region_id = region_file_id.region_id(); + let file_id = region_file_id.file_id(); let mut metrics = Metrics::new(WriteType::Flush); // Create index key for the SST file @@ -192,7 +193,8 @@ impl WriteCache { let store = self.file_cache.local_store(); let cleaner = TempFileCleaner::new(region_id, store.clone()); let write_res = store - .write(&cache_path, data.clone()) + .write_with(&cache_path, data.clone()) + .chunk(write_buffer_size.as_bytes() as usize) .await .context(crate::error::OpenDalSnafu); if let Err(e) = write_res { @@ -204,7 +206,6 @@ impl WriteCache { // Upload to remote store let upload_start = Instant::now(); - let region_file_id = RegionFileId::new(region_id, file_id); let remote_path = upload_request .dest_path_provider .build_sst_file_path(region_file_id); @@ -214,8 +215,11 @@ impl WriteCache { parquet_key, &remote_path, &upload_request.remote_store, - // `put_and_upload_sst` is only used by the flush path. - OperationType::Flush, + UploadOptions { + // `put_and_upload_sst` is only used by the flush path. + op_type: OperationType::Flush, + write_buffer_size, + }, ) .await { @@ -312,7 +316,15 @@ impl WriteCache { .build_sst_file_path(RegionFileId::new(region_id, sst.file_id)); let start = Instant::now(); if let Err(e) = self - .upload(parquet_key, &parquet_path, remote_store, op_type) + .upload( + parquet_key, + &parquet_path, + remote_store, + UploadOptions { + op_type, + write_buffer_size: write_opts.write_buffer_size, + }, + ) .await { err = Some(e); @@ -328,7 +340,15 @@ impl WriteCache { .build_index_file_path(RegionFileId::new(region_id, sst.file_id)); let start = Instant::now(); if let Err(e) = self - .upload(puffin_key, &puffin_path, remote_store, op_type) + .upload( + puffin_key, + &puffin_path, + remote_store, + UploadOptions { + op_type, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) .await { err = Some(e); @@ -399,7 +419,7 @@ impl WriteCache { index_key: IndexKey, upload_path: &str, remote_store: &ObjectStore, - op_type: OperationType, + options: UploadOptions, ) -> Result<()> { let region_id = index_key.region_id; let file_id = index_key.file_id; @@ -425,11 +445,11 @@ impl WriteCache { let upload_store = self.upload_store_wrapper.as_ref().map_or_else( || remote_store.clone(), - |wrapper| wrapper.wrap(remote_store.clone(), op_type), + |wrapper| wrapper.wrap(remote_store.clone(), options.op_type), ); let mut writer = upload_store .writer_with(upload_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(options.write_buffer_size.as_bytes() as usize) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .context(error::OpenDalSnafu)? @@ -486,6 +506,14 @@ pub struct SstUploadRequest { pub remote_store: ObjectStore, } +/// Policies for uploading a cached SST or index file. +pub(crate) struct UploadOptions { + /// Origin of the upload, used by the remote store wrapper. + pub op_type: OperationType, + /// Buffer size for the remote writer. + pub write_buffer_size: ReadableSize, +} + /// A structs to track files to upload and clean them if upload failed. pub(crate) struct UploadTracker { /// Id of the region to track. @@ -568,7 +596,7 @@ mod tests { use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - new_flat_source_from_record_batches, new_record_batch_by_range, + WriteChunkRecorder, new_flat_source_from_record_batches, new_record_batch_by_range, sst_file_handle_with_file_id, sst_region_metadata, }; @@ -617,7 +645,15 @@ mod tests { local_store.write(&cache_path, data.clone()).await.unwrap(); write_cache - .upload(key, upload_path, &original_store, OperationType::Compact) + .upload( + key, + upload_path, + &original_store, + UploadOptions { + op_type: OperationType::Compact, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) .await .unwrap(); @@ -633,16 +669,22 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] - async fn test_write_and_upload_sst() { + async fn test_write_and_upload_sst( + #[values(OperationType::Flush, OperationType::Compact)] op_type: OperationType, + ) { // TODO(QuenKar): maybe find a way to create some object server for testing, // and now just use local file system to mock. let mut env = TestEnv::new().await; - let mock_store = env.init_object_store_manager(); + let remote_chunks = WriteChunkRecorder::default(); + let mock_store = env.init_object_store_manager().layer(remote_chunks.layer()); let path_provider = RegionFilePathFactory::new("test".to_string(), PathType::Bare); let local_dir = create_temp_dir(""); - let local_store = new_fs_store(local_dir.path().to_str().unwrap()); + let local_chunks = WriteChunkRecorder::default(); + let local_store = + new_fs_store(local_dir.path().to_str().unwrap()).layer(local_chunks.layer()); let write_cache = env .create_write_cache(local_store.clone(), ReadableSize::mb(10)) @@ -658,7 +700,7 @@ mod tests { ]); let write_request = SstWriteRequest { - op_type: OperationType::Flush, + op_type, metadata, source, storage: None, @@ -681,6 +723,7 @@ mod tests { }; let write_opts = WriteOptions { + write_buffer_size: ReadableSize::kb(1), row_group_size: 512, ..Default::default() }; @@ -710,6 +753,14 @@ mod tests { .await .unwrap(); assert_eq!(remote_data.to_vec(), cache_data.to_vec()); + let chunk_size = write_opts.write_buffer_size.as_bytes() as usize; + assert!(remote_data.len() > chunk_size); + remote_chunks.assert_chunks(&sst_upload_path, chunk_size, remote_data.len()); + local_chunks.assert_chunks( + &write_cache.file_cache.cache_file_path(key), + chunk_size, + cache_data.len(), + ); // Check write cache contains the index key let index_key = IndexKey::new(region_id, file_id, FileType::Puffin(0)); @@ -721,6 +772,11 @@ mod tests { .await .unwrap(); assert_eq!(remote_index_data.to_vec(), cache_index_data.to_vec()); + remote_chunks.assert_chunks( + &index_upload_path, + DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize, + remote_index_data.len(), + ); // Removes the file from the cache. let sst_index_key = IndexKey::new(region_id, file_id, FileType::Parquet); diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 8ca4326e2b6..5164e2fd72b 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -164,6 +164,7 @@ pub struct MitoConfig { // Other configs: /// Buffer size for SST writing. + // TODO(#9240): Support per-object-store write buffer sizes for backend-specific tuning. pub sst_write_buffer_size: ReadableSize, /// Maximum number of SST files to scan concurrently (default 384). pub max_concurrent_scan_files: usize, diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 7b431c1b890..8a96276e5f7 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -63,7 +63,7 @@ use crate::request::{ }; use crate::schedule::CancellableTaskState; use crate::schedule::scheduler::{Job, SchedulerRef}; -use crate::sst::file::{FileMeta, UncommittedSsts}; +use crate::sst::file::{FileMeta, RegionFileId, UncommittedSsts}; use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, @@ -738,12 +738,18 @@ impl RegionFlushTask { let access_layer = self.access_layer.clone(); let cache_manager = self.cache_manager.clone(); let region_id = version.metadata.region_id; + let write_buffer_size = write_opts.write_buffer_size; let semaphore = self.flush_semaphore.clone(); let uncommitted = uncommitted.clone(); let task = common_runtime::spawn_global(async move { let _permit = semaphore.acquire().await.unwrap(); let metrics = access_layer - .put_sst(&encoded.data, region_id, &encoded.sst_info, &cache_manager) + .put_sst( + &encoded.data, + RegionFileId::new(region_id, encoded.sst_info.file_id), + &cache_manager, + write_buffer_size, + ) .await?; uncommitted.track(std::slice::from_ref(&encoded.sst_info)); FLUSH_FILE_TOTAL.inc(); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 67e8c4d8e1e..0c3ede4fb06 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -51,7 +51,7 @@ use vector_index::creator::VectorIndexer; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; -use crate::cache::write_cache::{UploadTracker, WriteCacheRef}; +use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; #[cfg(feature = "vector_index")] use crate::config::VectorIndexConfig; @@ -1022,7 +1022,10 @@ impl IndexBuildTask { puffin_key, &puffin_path, remote_store, - OperationType::Compact, + UploadOptions { + op_type: OperationType::Compact, + write_buffer_size: crate::sst::DEFAULT_WRITE_BUFFER_SIZE, + }, ) .await { diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 87eeb7bae5d..b49d2ba1e18 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -212,10 +212,10 @@ mod tests { }; use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - build_test_binary_test_region_metadata, new_flat_source_from_record_batches, - new_primary_key, new_record_batch_by_range, new_record_batch_with_custom_sequence, - new_sparse_primary_key, sst_file_handle, sst_file_handle_with_file_id, sst_region_metadata, - sst_region_metadata_with_encoding, + WriteChunkRecorder, build_test_binary_test_region_metadata, + new_flat_source_from_record_batches, new_primary_key, new_record_batch_by_range, + new_record_batch_with_custom_sequence, new_sparse_primary_key, sst_file_handle, + sst_file_handle_with_file_id, sst_region_metadata, sst_region_metadata_with_encoding, }; const FILE_DIR: &str = "/"; @@ -901,12 +901,14 @@ mod tests { .await; } + #[rstest::rstest] #[tokio::test] - async fn test_write_multiple_files() { + async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { common_telemetry::init_default_ut_logging(); // create test env let mut env = TestEnv::new().await; - let object_store = env.init_object_store_manager(); + let chunks = WriteChunkRecorder::default(); + let object_store = env.init_object_store_manager().layer(chunks.layer()); let metadata = Arc::new(sst_region_metadata()); let batches = vec![ new_record_batch_by_range(&["a", "a"], 0, 1000), @@ -921,6 +923,7 @@ mod tests { let source = new_flat_source_from_record_batches(batches); let write_opts = WriteOptions { + write_buffer_size: ReadableSize(write_buffer_size as u64), row_group_size: 50, max_file_size: Some(1024 * 16), ..Default::default() @@ -936,7 +939,7 @@ mod tests { metadata.clone(), IndexConfig::default(), NoopIndexBuilder, - path_provider, + path_provider.clone(), &mut metrics, ) .await; @@ -947,8 +950,18 @@ mod tests { .unwrap(); assert_eq!(2, files.len()); + // The configured buffer size must reach every writer, including split files. + assert_eq!(files.len(), chunks.num_files()); + let mut rows_read = 0; for f in &files { + assert!(f.file_size > write_buffer_size as u64); + chunks.assert_chunks( + &path_provider + .build_sst_file_path(RegionFileId::new(metadata.region_id, f.file_id)), + write_buffer_size, + f.file_size as usize, + ); let file_handle = sst_file_handle_with_file_id( f.file_id, f.time_range.0.value(), diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 9888b4ead53..b9ef15add34 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -65,8 +65,7 @@ use crate::sst::parquet::{ PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding, }; use crate::sst::{ - DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, - maybe_wrap_schema, + DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, }; /// Converts a flat RecordBatch for writing to parquet. @@ -181,7 +180,11 @@ pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathPro pub trait WriterFactory { type Writer: AsyncWrite + Send + Unpin; - fn create(&mut self, file_path: &str) -> impl Future>; + fn create( + &mut self, + file_path: &str, + write_buffer_size: usize, + ) -> impl Future>; } pub struct ObjectStoreWriterFactory { @@ -191,10 +194,10 @@ pub struct ObjectStoreWriterFactory { impl WriterFactory for ObjectStoreWriterFactory { type Writer = Compat; - async fn create(&mut self, file_path: &str) -> Result { + async fn create(&mut self, file_path: &str, write_buffer_size: usize) -> Result { self.object_store .writer_with(file_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(write_buffer_size) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .map(|v| v.into_futures_async_write().compat_write()) @@ -566,7 +569,9 @@ where self.current_file, )); let writer = SizeAwareWriter::new( - self.writer_factory.create(&sst_file_path).await?, + self.writer_factory + .create(&sst_file_path, opts.write_buffer_size.as_bytes() as usize) + .await?, self.bytes_written.clone(), ); let arrow_writer = diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs index 3f635c6e949..fd6321b0895 100644 --- a/src/mito2/src/test_util/sst_util.rs +++ b/src/mito2/src/test_util/sst_util.rs @@ -14,7 +14,8 @@ //! Utilities for testing SSTs. -use std::sync::Arc; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use api::v1::{OpType, SemanticType}; use arrow_schema::Schema; @@ -28,6 +29,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::ValueRef; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use object_store::layers::mock::{self, MockLayer, MockLayerBuilder, oio}; use store_api::metadata::{ ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef, }; @@ -45,6 +47,73 @@ use crate::test_util::{new_batch_builder, new_noop_file_purger}; /// Test region id. const REGION_ID: RegionId = RegionId::new(0, 0); +/// Records write chunks at the object-store boundary without replacing filesystem I/O. +#[derive(Default)] +pub(crate) struct WriteChunkRecorder { + chunk_sizes: Arc>>>, +} + +impl WriteChunkRecorder { + pub(crate) fn layer(&self) -> MockLayer { + let chunk_sizes = self.chunk_sizes.clone(); + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + Box::new(ChunkRecordingWriter { + inner, + path: path.to_string(), + chunk_sizes: chunk_sizes.clone(), + }) + })) + .build() + .unwrap() + } + + pub(crate) fn num_files(&self) -> usize { + self.chunk_sizes.lock().unwrap().len() + } + + pub(crate) fn assert_chunks(&self, path: &str, chunk_size: usize, file_size: usize) { + let mut expected = vec![chunk_size; file_size / chunk_size]; + let remainder = file_size % chunk_size; + if remainder > 0 { + expected.push(remainder); + } + assert_eq!( + self.chunk_sizes.lock().unwrap().get(path), + Some(&expected), + "unexpected write chunks for {path}" + ); + } +} + +struct ChunkRecordingWriter { + inner: oio::Writer, + path: String, + chunk_sizes: Arc>>>, +} + +impl oio::Write for ChunkRecordingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + let size = buffer.len(); + self.inner.write(buffer).await?; + self.chunk_sizes + .lock() + .unwrap() + .entry(self.path.clone()) + .or_default() + .push(size); + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + self.inner.close().await + } + + async fn abort(&mut self) -> mock::Result<()> { + self.inner.abort().await + } +} + /// Creates a new region metadata for testing SSTs with specified encoding. /// /// Dense schema: tag_0, tag_1, field_0, ts diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index f9b367c3dce..412f3563298 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -27,7 +27,7 @@ derive_builder = { workspace = true, optional = true } futures.workspace = true humantime-serde.workspace = true lazy_static.workspace = true -opendal = { version = "0.58", features = [ +opendal = { version = "0.59.2", features = [ "layers-tracing", "layers-prometheus", "services-azblob", @@ -37,7 +37,7 @@ opendal = { version = "0.58", features = [ "services-oss", "services-s3", ] } -opendal-http-transport-reqwest = "0.58" +opendal-http-transport-reqwest = "0.59.2" prometheus.workspace = true reqwest.workspace = true serde.workspace = true diff --git a/src/object-store/src/layers/mock.rs b/src/object-store/src/layers/mock.rs index 6b919aabb69..270f776f935 100644 --- a/src/object-store/src/layers/mock.rs +++ b/src/object-store/src/layers/mock.rs @@ -96,6 +96,7 @@ impl Service for MockService { type Lister = oio::Lister; type Deleter = oio::Deleter; type Copier = oio::Copier; + type Composer = oio::Composer; fn info(&self) -> opendal::raw::ServiceInfo { self.inner.info() @@ -165,7 +166,6 @@ impl Service for MockService { from: &str, to: &str, args: OpCopy, - opts: opendal::raw::OpCopier, ) -> Result { if let Some(result) = self .copy_interceptor @@ -176,7 +176,16 @@ impl Service for MockService { return Ok(Box::new(oio::OneShotCopier::completed()) as oio::Copier); } - self.inner.copy(ctx, from, to, args, opts) + self.inner.copy(ctx, from, to, args) + } + + fn compose( + &self, + ctx: &OperationContext, + to: &str, + args: opendal::raw::OpCompose, + ) -> Result { + self.inner.compose(ctx, to, args) } async fn rename( diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index b940755c5b3..18afa3ac820 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -24,8 +24,8 @@ use cap_std::fs::{Dir, DirEntry, OpenOptions, ReadDir}; use opendal::layers::SimulateLayer; use opendal::raw::*; use opendal::{ - Buffer, BytesRange, Capability, EntryMode, Error, ErrorKind, Metadata, OperationContext, - Operator, Result, + Buffer, BytesRange, Capability, EntryMode, Error, ErrorKind, Metadata, MetadataBuilder, + OperationContext, Operator, Result, }; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; @@ -165,19 +165,18 @@ fn parse_write_error(error: io::Error, if_not_exists: bool) -> Error { } fn metadata_from_fs(metadata: cap_std::fs::Metadata) -> Result { - let mode = if metadata.is_dir() { - EntryMode::DIR + let mut builder = if metadata.is_dir() { + MetadataBuilder::dir() } else if metadata.is_file() { - EntryMode::FILE + MetadataBuilder::file(metadata.len()) } else { - EntryMode::Unknown + MetadataBuilder::unknown() }; - Ok(Metadata::new(mode) - .with_content_length(metadata.len()) - .with_last_modified(Timestamp::try_from( - metadata.modified().map_err(new_std_io_error)?.into_std(), - )?)) + builder.last_modified(Timestamp::try_from( + metadata.modified().map_err(new_std_io_error)?.into_std(), + )?); + Ok(builder.build()) } #[derive(Clone, Debug)] @@ -219,6 +218,7 @@ impl Service for SecureFsBackend { type Lister = SecureFsLister; type Deleter = oio::OneShotDeleter; type Copier = (); + type Composer = (); fn info(&self) -> ServiceInfo { self.info.clone() @@ -301,14 +301,7 @@ impl Service for SecureFsBackend { }) } - fn copy( - &self, - _: &OperationContext, - _: &str, - _: &str, - _: OpCopy, - _: OpCopier, - ) -> Result { + fn copy(&self, _: &OperationContext, _: &str, _: &str, _: OpCopy) -> Result { Err(Error::new( ErrorKind::Unsupported, "operation is not supported", @@ -449,11 +442,11 @@ impl oio::Write for SecureFsWriter { file.flush().await.map_err(new_std_io_error)?; file.sync_all().await.map_err(new_std_io_error)?; let metadata = file.metadata().await.map_err(new_std_io_error)?; - Ok(Metadata::new(EntryMode::FILE) - .with_content_length(metadata.len()) - .with_last_modified(Timestamp::try_from( - metadata.modified().map_err(new_std_io_error)?, - )?)) + let mut builder = MetadataBuilder::file(metadata.len()); + builder.last_modified(Timestamp::try_from( + metadata.modified().map_err(new_std_io_error)?, + )?); + Ok(builder.build()) } async fn abort(&mut self) -> Result<()> { @@ -519,7 +512,7 @@ impl oio::List for SecureFsLister { } else { &display_prefix }, - Metadata::new(EntryMode::DIR), + MetadataBuilder::dir().build(), ); self.entries = vec![current_path].into_iter(); } @@ -589,7 +582,7 @@ fn read_list_entry(entry: DirEntry, display_prefix: &str) -> io::Result match metadata_from_fs(metadata) { From e27167139f861994f2a62aa24b62259d2c58a5f7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 21 Sep 2026 07:56:28 +0000 Subject: [PATCH 40/63] refactor: isolate logical batch scheduling and flushing (#9211) Signed-off-by: WenyXu --- src/servers/src/batcher/logical_table.rs | 1307 +---------------- .../src/batcher/logical_table/batch.rs | 744 ++++++++++ .../batcher/logical_table/pending_worker.rs | 656 +++++++++ 3 files changed, 1418 insertions(+), 1289 deletions(-) create mode 100644 src/servers/src/batcher/logical_table/batch.rs create mode 100644 src/servers/src/batcher/logical_table/pending_worker.rs diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index 9543b0d3a88..8e7bac8eb2f 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod batch; mod batch_convert; +mod pending_worker; mod region_write; #[cfg(test)] mod test_util; @@ -20,7 +22,7 @@ mod test_util; use std::collections::{HashMap, HashSet}; use std::num::NonZeroUsize; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::{ColumnSchema, RowInsertRequests, Rows}; @@ -28,17 +30,15 @@ use arrow::datatypes::Schema as ArrowSchema; use async_trait::async_trait; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; -use common_batcher::flush_policy::FlushTrigger; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::notifier::{Notifier, run_notifier}; -use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_runtime::spawn_global; -use common_telemetry::{debug, error, warn}; +use common_telemetry::{error, warn}; use datatypes::timestamp::append_timestamps; use meter_core::data::MeterRecord; use meter_macros::write_meter; @@ -46,16 +46,12 @@ use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; use table::metadata::TableId; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, mpsc, oneshot}; +use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; +pub use crate::batcher::logical_table::batch::flush_batch_physical; pub use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; -use crate::batcher::logical_table::batch_convert::{ - concat_modified_batches, transform_logical_batches_to_physical, -}; -use crate::batcher::logical_table::region_write::{ - CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter, - PartitionManagerPhysicalFlushAdapter, encode_region_write_requests, - flush_region_writes_concurrently, plan_region_batches, resolve_region_targets, +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, @@ -64,9 +60,7 @@ pub use crate::batcher::logical_table::region_write::{ use crate::error; use crate::error::{Error, Result}; use crate::metrics::{ - FLOW_NOTIFICATION_DROPPED, FLUSH_DROPPED_ROWS, FLUSH_ELAPSED, FLUSH_FAILURES, FLUSH_ROWS, - FLUSH_TOTAL, PENDING_BATCHES, PENDING_ROWS, PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED, - PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED, PENDING_WORKERS, + FLOW_NOTIFICATION_DROPPED, PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED, PENDING_WORKERS, }; use crate::prom_row_builder::{ build_prom_create_table_schema_from_proto, identify_missing_columns_from_proto, @@ -99,7 +93,7 @@ const MAX_CONCURRENT_FLOW_NOTIFICATIONS: NonZeroUsize = NonZeroUsize::new(8).unw #[async_trait] pub trait PendingRowsSchemaAlterer: Send + Sync { - /// FlushBatch-create multiple logical tables that are missing. + /// Batch-create multiple logical tables that are missing. /// Each entry is `(table_name, request_schema)`. async fn create_tables_if_missing_batch( &self, @@ -110,7 +104,7 @@ pub trait PendingRowsSchemaAlterer: Send + Sync { ctx: QueryContextRef, ) -> Result<()>; - /// FlushBatch-alter multiple logical tables to add missing tag columns. + /// Batch-alter multiple logical tables to add missing tag columns. /// Each entry is `(table_name, missing_column_names)`. async fn add_missing_prom_tag_columns_batch( &self, @@ -142,42 +136,6 @@ struct TableResolutionPlan { tables_to_alter: Vec<(String, Vec)>, } -struct PendingBatch { - tables: HashMap, - db_string: String, - ctx: QueryContextRef, -} - -struct FlushWaiter { - response_tx: oneshot::Sender>>, - _permit: Arc, -} - -struct FlushBatch { - table_batches: Vec, - total_row_count: usize, - db_string: String, - ctx: QueryContextRef, - waiters: Vec, -} - -#[derive(Clone)] -struct PendingWorker { - tx: mpsc::Sender, -} - -enum WorkerCommand { - Submit { - table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, - total_rows: usize, - ctx: QueryContextRef, - response_tx: oneshot::Sender>>, - _permit: Arc, - }, - #[cfg(test)] - Ack { ack_tx: oneshot::Sender<()> }, -} - // Requests can share a batch only when their write target and WAL policy match. fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { let physical_table = ctx @@ -518,7 +476,7 @@ impl LogicalTablePendingRowsBatcher { } impl LogicalTablePendingRowsBatcher { - /// FlushBatch-creates missing tables, refreshes their schema metadata, and + /// Batch-creates missing tables, refreshes their schema metadata, and /// enqueues follow-up alters for extra tag columns discovered in later rows. async fn create_missing_tables_and_refresh_schemas( &self, @@ -627,7 +585,7 @@ impl LogicalTablePendingRowsBatcher { } impl LogicalTablePendingRowsBatcher { - /// FlushBatch-alters tables that have missing tag columns and refreshes the + /// Batch-alters tables that have missing tag columns and refreshes the /// in-memory schema map used for row alignment. async fn alter_tables_and_refresh_schemas( &self, @@ -763,325 +721,6 @@ impl Drop for LogicalTablePendingRowsBatcher { } } -impl PendingBatch { - fn new(ctx: QueryContextRef) -> Self { - let db_string = ctx.get_db_string(); - Self { - tables: HashMap::new(), - db_string, - ctx, - } - } - - fn add_table_batch( - &mut self, - table_name: String, - table_id: TableId, - record_batch: RecordBatchWithTsIdx, - ) { - let entry = self.tables.entry(table_id).or_insert_with(|| TableBatch { - table_name, - table_id, - batches: Vec::new(), - row_count: 0, - }); - entry.row_count += record_batch.batch.num_rows(); - entry.batches.push(record_batch); - } -} - -#[allow(clippy::too_many_arguments)] -fn start_worker( - key: BatchKey, - worker_tx: mpsc::Sender, - workers: Arc>, - mut rx: mpsc::Receiver, - shutdown: broadcast::Sender<()>, - partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, - catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, - worker_idle_timeout: Duration, - flush_policy: TimingFlushPolicy, - flush_limiter: FlushLimiter, -) { - tokio::spawn(async move { - // The business batch and pending flush are populated and drained together. - let mut batch = None; - let mut pending_flush = PendingCore::new(flush_policy); - let mut shutdown_rx = shutdown.subscribe(); - let idle_deadline = tokio::time::Instant::now() + worker_idle_timeout; - let idle_timer = tokio::time::sleep_until(idle_deadline); - tokio::pin!(idle_timer); - - loop { - tokio::select! { - cmd = rx.recv() => { - match cmd { - Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit }) => { - let submitted_at = tokio::time::Instant::now(); - idle_timer.as_mut().reset(submitted_at + worker_idle_timeout); - - pending_flush.submit( - FlushWaiter { response_tx, _permit }, - total_rows, - ); - let pending_batch = batch.get_or_insert_with(||{ - PENDING_BATCHES.inc(); - PendingBatch::new(ctx) - }); - - for (table_name, table_id, record_batch) in table_batches { - pending_batch.add_table_batch(table_name, table_id, record_batch); - } - - PENDING_ROWS.add(total_rows as i64); - - if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) { - spawn_flush( - flush, - partition_manager.clone(), - node_manager.clone(), - catalog_manager.clone(), - flow_notification_tx.clone(), - flush_limiter.clone(), - ).await; - } - } - None => { - if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { - flush_batch_with_managers( - flush, - partition_manager.clone(), - node_manager.clone(), - catalog_manager.clone(), - flow_notification_tx.clone(), - ).await; - } - break; - } - #[cfg(test)] - Some(WorkerCommand::Ack { ack_tx }) => { - let _ = ack_tx.send(()); - } - } - } - _ = &mut idle_timer => { - if !should_close_worker_on_idle_timeout( - pending_flush.total_rows(), - rx.len(), - ) { - idle_timer - .as_mut() - .reset(tokio::time::Instant::now() + worker_idle_timeout); - continue; - } - - debug!( - "Closing idle pending rows worker due to timeout: catalog={}, schema={}, physical_table={}", - key.catalog, - key.schema, - key.physical_table - ); - break; - } - _ = pending_flush.wait_flush() => { - if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) { - spawn_flush( - flush, - partition_manager.clone(), - node_manager.clone(), - catalog_manager.clone(), - flow_notification_tx.clone(), - flush_limiter.clone(), - ).await; - } - } - _ = shutdown_rx.recv() => { - if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { - flush_batch_with_managers( - flush, - partition_manager.clone(), - node_manager.clone(), - catalog_manager.clone(), - flow_notification_tx.clone(), - ).await; - } - break; - } - } - } - - remove_worker_if_same_channel(workers.as_ref(), &key, &worker_tx).await; - }); -} - -async fn remove_worker_if_same_channel( - workers: &WorkerRegistry, - key: &BatchKey, - worker_tx: &mpsc::Sender, -) -> bool { - if workers.remove_if_same(key, worker_tx).await { - PENDING_WORKERS.set(workers.len().await as i64); - true - } else { - false - } -} - -fn should_close_worker_on_idle_timeout(total_row_count: usize, queued_requests: usize) -> bool { - total_row_count == 0 && queued_requests == 0 -} - -/// Transfers the ready batch, or drains unconditionally when no trigger is given. -/// Execution and flush permits remain owned by the caller. -fn drain_batch( - batch: &mut Option, - pending_flush: &mut PendingCore, - trigger: Option, -) -> Option { - let total_row_count = pending_flush.total_rows(); - let waiters = match trigger { - Some(trigger) => pending_flush.take_ready(trigger)?, - None => pending_flush.take_pending()?, - }; - let batch = batch.take()?; - - if total_row_count == 0 { - return None; - } - - let table_batches = batch.tables.into_values().collect(); - - PENDING_ROWS.sub(total_row_count as i64); - PENDING_BATCHES.dec(); - - Some(FlushBatch { - table_batches, - total_row_count, - db_string: batch.db_string, - ctx: batch.ctx, - waiters, - }) -} - -async fn spawn_flush( - flush: FlushBatch, - partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, - catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, - flush_limiter: FlushLimiter, -) { - match flush_limiter.acquire().await { - Ok(permit) => { - tokio::spawn(async move { - let _permit = permit; - flush_batch_with_managers( - flush, - partition_manager, - node_manager, - catalog_manager, - flow_notification_tx, - ) - .await; - }); - } - Err(err) => { - warn!(err; "Flush semaphore closed, flushing inline"); - flush_batch_with_managers( - flush, - partition_manager, - node_manager, - catalog_manager, - flow_notification_tx, - ) - .await; - } - } -} - -async fn flush_batch_with_managers( - flush: FlushBatch, - partition_manager: PartitionRuleManagerRef, - node_manager: NodeManagerRef, - catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, -) { - let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; - let node_requester = NodeManagerPhysicalFlushAdapter { - node_manager: node_manager.clone(), - }; - let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager }; - flush_batch( - flush, - &partition_provider, - &node_requester, - &catalog_provider, - flow_notification_tx, - ) - .await; -} - -async fn flush_batch( - flush: FlushBatch, - partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), - catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), - flow_notification_tx: Notifier, -) { - let FlushBatch { - table_batches, - total_row_count, - db_string, - ctx, - waiters, - } = flush; - let start = Instant::now(); - - // 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) - .to_string(); - let result = flush_batch_physical( - &table_batches, - &physical_table_name, - &ctx, - partition_manager, - node_manager, - catalog_manager, - ) - .await; - - let elapsed = start.elapsed().as_secs_f64(); - FLUSH_ELAPSED.observe(elapsed); - - debug!( - "Pending rows batch flushed, total rows: {}, elapsed time: {}s", - total_row_count, elapsed - ); - - match result { - Ok(affected_rows) => { - FLUSH_TOTAL.inc(); - 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); - } - Err(err) => { - FLUSH_FAILURES.inc(); - FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64); - notify_waiters(waiters, Err(err)); - } - } -} - fn extract_timestamps(table_batch: &TableBatch) -> Vec { let mut timestamps = Vec::with_capacity(table_batch.row_count); for batch in &table_batch.batches { @@ -1217,113 +856,13 @@ fn notify_flow_dirty_windows_after_flush( enqueue_flow_notifications(table_batches, &tx); } -/// Flushes a batch of logical table rows by transforming them into the physical table format -/// and writing them to the appropriate datanode regions. -/// -/// This function performs the end-to-end physical flush pipeline: -/// 1. Resolves the physical table metadata and column ID mapping. -/// 2. Fetches the physical table's partition rule. -/// 3. Transforms each logical table batch into the physical (sparse primary key) format. -/// 4. Concatenates all transformed batches into a single combined batch. -/// 5. Splits the combined batch by partition rule and sends region write requests -/// concurrently to the target datanodes. -pub async fn flush_batch_physical( - table_batches: &[TableBatch], - physical_table_name: &str, - ctx: &QueryContextRef, - partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), - node_manager: &(impl PhysicalFlushNodeRequester + ?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 - .with_label_values(&["flush_physical_resolve_table"]) - .start_timer(); - catalog_manager - .physical_table( - ctx.current_catalog(), - &ctx.current_schema(), - physical_table_name, - ctx.as_ref(), - ) - .await? - .with_context(|| error::InternalSnafu { - err_msg: format!( - "Physical table '{}' not found during pending flush", - physical_table_name - ), - })? - }; - - let physical_table_info = physical_table.table_info; - let name_to_ids = physical_table - .col_name_to_ids - .with_context(|| error::InternalSnafu { - err_msg: format!( - "Physical table '{}' has no column IDs for pending flush", - physical_table_name - ), - })?; - - // 2. Get the physical table's partition rule (one lookup instead of N) - let partition_rule = { - let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_fetch_partition_rule"]) - .start_timer(); - partition_manager - .find_table_partition_rule(physical_table_info.as_ref()) - .await? - }; - let partition_columns = partition_rule.partition_columns(); - let partition_columns_set: HashSet<&str> = - 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)?; - - // 4. Concatenate all modified batches (all share the same physical schema) - let combined_batch = 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( - 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 -} - -fn notify_waiters(waiters: Vec, result: Result<()>) { - let shared_result = result.map_err(Arc::new); - for waiter in waiters { - let _ = waiter.response_tx.send(match &shared_result { - Ok(()) => Ok(()), - Err(error) => Err(Arc::clone(error)), - }); - // waiter._permit is dropped here, releasing the inflight semaphore slot - } -} - #[cfg(test)] mod tests { use std::any::Any; - use std::collections::HashMap; - use std::future::poll_fn; - use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; - use std::task::Poll; use std::time::Duration; - use api::region::RegionResponse; use api::v1::meta::Peer; - use api::v1::region::RegionRequest; use api::v1::value::ValueData; use api::v1::{ ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, @@ -1333,21 +872,10 @@ mod tests { use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; - use catalog::error::Result as CatalogResult; - use catalog::memory::MemoryCatalogManager; - use common_batcher::flush_limiter::FlushLimiter; - use common_batcher::flush_policy::FlushTrigger; - use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::notifier::Notifier; - use common_batcher::pending_worker::PendingWorker as PendingCore; - use common_batcher::request_limiter::RequestLimiter; - use common_batcher::worker_registry::WorkerRegistry; - use common_meta::cache::{ - TableFlownodeSetCacheRef, new_table_flownode_set_cache, new_table_route_cache, - }; + use common_meta::cache::new_table_flownode_set_cache; use common_meta::error::Result as MetaResult; use common_meta::instruction::{CacheIdent, CreateFlow}; - use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackend, TxnService}; use common_meta::node_manager::NodeManagerRef; use common_meta::rpc::store::{ @@ -1356,36 +884,18 @@ mod tests { PutResponse, RangeRequest, RangeResponse, }; use common_query::prelude::greptime_timestamp; - use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use moka::future::CacheBuilder; - use partition::cache::new_partition_info_cache; - use partition::error::Result as PartitionResult; - use partition::manager::PartitionRuleManager; - 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::{Notify, Semaphore, broadcast, mpsc, oneshot}; - use tokio::time::advance; + use tokio::sync::{Notify, mpsc, oneshot}; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; - use crate::batcher::logical_table::region_write::{ - PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, - PhysicalTableMetadata, - }; use crate::batcher::logical_table::test_util::{ - ConcurrentMockNodeManager, FlowNotificationMockNodeManager, RecordingFlownode, - mock_aligned_tag_batch, mock_table_flownode_cache, mock_tag_batch, + FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch, + mock_table_flownode_cache, mock_tag_batch, }; use crate::batcher::logical_table::{ - BatchKey, FlushBatch, FlushWaiter, LogicalTablePendingRowsBatcher, PendingBatch, - WorkerCommand, drain_batch, extract_timestamps, flush_batch, flush_batch_physical, - notify_flow_dirty_windows_after_flush, notify_waiters, remove_worker_if_same_channel, - should_close_worker_on_idle_timeout, start_flow_notification_worker, start_worker, + LogicalTablePendingRowsBatcher, extract_timestamps, notify_flow_dirty_windows_after_flush, try_enqueue_flow_notification, }; - use crate::error; - use crate::error::Error; use crate::metrics::FLOW_NOTIFICATION_DROPPED; use crate::prom_row_builder::rows_to_aligned_record_batch; @@ -1600,138 +1110,6 @@ mod tests { assert_eq!(closed_before + 1, closed.get()); } - fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata { - let schema = Arc::new( - DtSchema::try_new(vec![ - DtColumnSchema::new( - "__primary_key", - datatypes::prelude::ConcreteDataType::binary_datatype(), - false, - ), - DtColumnSchema::new( - "greptime_timestamp", - datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - DtColumnSchema::new( - "greptime_value", - datatypes::prelude::ConcreteDataType::float64_datatype(), - true, - ), - DtColumnSchema::new( - "tag1", - datatypes::prelude::ConcreteDataType::string_datatype(), - true, - ), - ]) - .unwrap(), - ); - let mut table_info = test_table_info(table_id, "phy", "public", "greptime", schema); - table_info.meta.column_ids = vec![0, 1, 2, 3]; - - PhysicalTableMetadata { - table_info: Arc::new(table_info), - col_name_to_ids: Some(HashMap::from([("tag1".to_string(), 3)])), - } - } - - struct MockFlushCatalogProvider { - table: Option, - } - - #[async_trait] - impl PhysicalFlushCatalogProvider for MockFlushCatalogProvider { - async fn physical_table( - &self, - _catalog: &str, - _schema: &str, - _table_name: &str, - _query_ctx: &session::context::QueryContext, - ) -> CatalogResult> { - Ok(self.table.clone()) - } - } - - struct SingleRegionPartitionRule; - - impl PartitionRule for SingleRegionPartitionRule { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn partition_columns(&self) -> &[String] { - &[] - } - - fn find_region( - &self, - _values: &[datatypes::prelude::Value], - ) -> partition::error::Result { - unimplemented!() - } - - fn split_record_batch( - &self, - record_batch: &RecordBatch, - ) -> partition::error::Result> - { - Ok(HashMap::from([( - 1, - RegionMask::new( - arrow::array::BooleanArray::from(vec![true; record_batch.num_rows()]), - record_batch.num_rows(), - ), - )])) - } - } - - struct MockFlushPartitionProvider { - partition_rule_calls: Arc, - region_leader_calls: Arc, - } - - #[async_trait] - impl PhysicalFlushPartitionProvider for MockFlushPartitionProvider { - async fn find_table_partition_rule( - &self, - _table_info: &table::metadata::TableInfo, - ) -> PartitionResult { - 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(), - }) - } - } - - #[derive(Default)] - struct MockFlushNodeRequester { - writes: Arc, - fail: bool, - } - - #[async_trait] - impl PhysicalFlushNodeRequester for MockFlushNodeRequester { - async fn handle( - &self, - _peer: &Peer, - _request: RegionRequest, - ) -> 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)) - } - } - #[test] fn test_collect_non_empty_table_rows_filters_empty_payloads() { let requests = RowInsertRequests { @@ -1760,133 +1138,6 @@ mod tests { assert_eq!(2, table_rows[0].1.rows.len()); } - #[tokio::test] - async fn test_drain_batch_takes_initialized_pending_batch_from_option() { - let ctx = session::context::QueryContext::arc(); - let (response_tx, _response_rx) = oneshot::channel(); - let permit = Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap(); - let mut pending_flush = - PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 1).unwrap()); - pending_flush.submit( - FlushWaiter { - response_tx, - _permit: Arc::new(permit), - }, - 1, - ); - let mut batch = Some(PendingBatch { - tables: HashMap::from([( - 42, - TableBatch { - table_name: "cpu".to_string(), - table_id: 42, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }, - )]), - db_string: ctx.get_db_string(), - ctx: ctx.clone(), - }); - - let flush = drain_batch( - &mut batch, - &mut pending_flush, - Some(FlushTrigger::Submission), - ) - .unwrap(); - - assert!(batch.is_none()); - assert!(pending_flush.is_empty()); - assert_eq!(0, pending_flush.total_rows()); - assert_eq!(1, flush.waiters.len()); - assert_eq!(1, flush.total_row_count); - assert_eq!(1, flush.table_batches.len()); - assert_eq!(ctx.get_db_string(), flush.db_string); - assert_eq!(ctx.current_catalog(), flush.ctx.current_catalog()); - } - - #[tokio::test] - async fn test_drain_batch_preserves_unready_state_and_clears_zero_rows() { - for total_rows in [0, 1] { - let mut pending_flush = - PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 2).unwrap()); - let mut batch = Some(PendingBatch::new(session::context::QueryContext::arc())); - let semaphore = Arc::new(Semaphore::new(1)); - let (response_tx, mut response_rx) = oneshot::channel(); - pending_flush.submit( - FlushWaiter { - response_tx, - _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()), - }, - total_rows, - ); - assert!( - drain_batch( - &mut batch, - &mut pending_flush, - Some(FlushTrigger::Submission) - ) - .is_none() - ); - assert!(batch.is_some()); - assert!(!pending_flush.is_empty()); - assert_eq!(total_rows, pending_flush.total_rows()); - assert_eq!(0, semaphore.available_permits()); - - let drained = drain_batch(&mut batch, &mut pending_flush, None); - assert!(batch.is_none()); - assert!(pending_flush.is_empty()); - assert_eq!(0, pending_flush.total_rows()); - if total_rows == 0 { - assert!(drained.is_none()); - assert_eq!(1, semaphore.available_permits()); - assert!(matches!( - response_rx.try_recv(), - Err(oneshot::error::TryRecvError::Closed) - )); - } else { - let drained = drained.unwrap(); - assert_eq!(1, drained.total_row_count); - assert_eq!(1, drained.waiters.len()); - assert_eq!(0, semaphore.available_permits()); - drop(drained); - assert_eq!(1, semaphore.available_permits()); - } - } - } - - #[test] - fn test_pending_batch_keeps_same_name_batches_with_distinct_table_ids() { - let ctx = session::context::QueryContext::arc(); - let mut pending_batch = PendingBatch::new(ctx); - - pending_batch.add_table_batch( - "cpu".to_string(), - 42, - mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0), - ); - pending_batch.add_table_batch( - "cpu".to_string(), - 43, - mock_aligned_tag_batch("tag1", "host-1", 2000, 2.0), - ); - - assert_eq!(2, pending_batch.tables.len()); - assert_eq!(42, pending_batch.tables[&42].table_id); - assert_eq!(43, pending_batch.tables[&43].table_id); - assert_eq!("cpu", pending_batch.tables[&42].table_name); - assert_eq!("cpu", pending_batch.tables[&43].table_name); - } - - fn mock_flow_notification_sender( - cache: TableFlownodeSetCacheRef, - node_manager: NodeManagerRef, - ) -> Notifier { - let (tx, rx) = Notifier::try_new(16).unwrap(); - start_flow_notification_worker(rx, cache, node_manager); - tx - } - struct BlockingRangeKvBackend { range_started: Mutex>>, range_release: Arc, @@ -2084,528 +1335,6 @@ mod tests { ); } - #[tokio::test] - async fn test_flush_batch_notifies_flownode_after_successful_physical_write() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); - let ctx = session::context::QueryContext::arc(); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let writes = Arc::new(AtomicUsize::new(0)); - - flush_batch( - FlushBatch { - table_batches, - total_row_count: 1, - db_string: ctx.get_db_string(), - ctx, - waiters: Vec::new(), - }, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &MockFlushNodeRequester { - writes: writes.clone(), - fail: false, - }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - flow_notification_tx, - ) - .await; - - assert_eq!(1, writes.load(Ordering::SeqCst)); - let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) - .await - .unwrap() - .unwrap(); - assert_eq!( - vec![api::v1::flow::DirtyWindowRequest { - table_id, - timestamps: vec![1000], - time_ranges: Vec::new(), - }], - requests.requests - ); - } - - #[tokio::test] - async fn test_flush_batch_does_not_notify_flownode_after_physical_write_error() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); - let ctx = session::context::QueryContext::arc(); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let writes = Arc::new(AtomicUsize::new(0)); - - flush_batch( - FlushBatch { - table_batches, - total_row_count: 1, - db_string: ctx.get_db_string(), - ctx, - waiters: Vec::new(), - }, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &MockFlushNodeRequester { - writes: writes.clone(), - fail: true, - }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - flow_notification_tx, - ) - .await; - - assert_eq!(1, writes.load(Ordering::SeqCst)); - assert!( - tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) - .await - .is_err() - ); - } - - #[tokio::test] - async fn test_cancelled_waiter_retains_request_slot_until_notification() { - let limiter = RequestLimiter::try_new(1).unwrap(); - let (response_tx, response_rx) = oneshot::channel(); - let waiter = FlushWaiter { - response_tx, - _permit: limiter.acquire().await.unwrap(), - }; - drop(response_rx); - let next = limiter.acquire(); - tokio::pin!(next); - assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await); - notify_waiters(vec![waiter], Ok(())); - let _permit = next.await.unwrap(); - } - - #[tokio::test] - async fn test_remove_worker_if_same_channel_removes_matching_entry() { - let workers = WorkerRegistry::new(); - let key = BatchKey { - catalog: "greptime".to_string(), - schema: "public".to_string(), - physical_table: "phy".to_string(), - skip_wal: false, - }; - - let (tx, _rx) = mpsc::channel::(1); - workers.get_or_insert_with(key.clone(), || tx.clone()).await; - - assert!(remove_worker_if_same_channel(&workers, &key, &tx).await); - assert!(workers.is_empty().await); - } - - #[tokio::test] - async fn test_remove_worker_if_same_channel_keeps_newer_entry() { - let workers = WorkerRegistry::new(); - let key = BatchKey { - catalog: "greptime".to_string(), - schema: "public".to_string(), - physical_table: "phy".to_string(), - skip_wal: false, - }; - - let (stale_tx, _stale_rx) = mpsc::channel::(1); - let (fresh_tx, _fresh_rx) = mpsc::channel::(1); - workers - .get_or_insert_with(key.clone(), || fresh_tx.clone()) - .await; - - assert!(!remove_worker_if_same_channel(&workers, &key, &stale_tx).await); - assert!(workers.get(&key).await.is_some()); - assert!(workers.get(&key).await.unwrap().same_channel(&fresh_tx)); - } - - #[test] - fn test_worker_idle_timeout_close_decision() { - assert!(should_close_worker_on_idle_timeout(0, 0)); - assert!(!should_close_worker_on_idle_timeout(1, 0)); - assert!(!should_close_worker_on_idle_timeout(0, 1)); - } - - const WORKER_TEST_TIMEOUT: Duration = Duration::from_secs(30); - - async fn submit_mock_worker_batch( - worker_tx: &mpsc::Sender, - total_rows: usize, - timestamp: i64, - ) -> oneshot::Receiver>> { - let (response_tx, response_rx) = oneshot::channel(); - let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); - worker_tx - .send(WorkerCommand::Submit { - table_batches: vec![( - "cpu".to_string(), - 42, - mock_aligned_tag_batch("tag1", "host-1", timestamp, 1.0), - )], - total_rows, - ctx: session::context::QueryContext::arc(), - response_tx, - _permit: Arc::new(permit), - }) - .await - .unwrap(); - - // The channel is FIFO, so the ack proves the worker has dequeued and - // processed the submission (anchoring the flush deadline) before the - // caller advances virtual time. - let (ack_tx, ack_rx) = oneshot::channel(); - worker_tx.send(WorkerCommand::Ack { ack_tx }).await.unwrap(); - ack_rx - .await - .expect("worker exited before acking the submitted batch"); - - response_rx - } - - async fn receive_mock_flush_result( - response_rx: oneshot::Receiver>>, - context: &str, - ) -> std::result::Result<(), Arc> { - // Under paused time the timeout auto-advances the clock and fires - // deterministically if the flush never completes. - tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx) - .await - .unwrap_or_else(|_| panic!("{context}")) - .expect("flush result channel closed without a result") - } - - fn assert_missing_physical_table(result: std::result::Result<(), Arc>) { - let err = result.expect_err("the empty catalog should make the flush fail"); - assert!( - matches!( - err.as_ref(), - Error::Internal { err_msg } - if err_msg.contains("not found during pending flush") - ), - "unexpected flush error: {err}" - ); - } - - #[tokio::test(start_paused = true)] - async fn test_worker_preserves_first_deadline_and_inline_shutdown() { - let flush_interval = Duration::from_secs(10); - let worker_idle_timeout = Duration::from_secs(30); - let key = BatchKey { - catalog: "greptime".to_string(), - schema: "public".to_string(), - physical_table: "phy".to_string(), - skip_wal: false, - }; - let workers = Arc::new(WorkerRegistry::new()); - let (worker_tx, worker_rx) = mpsc::channel(1); - workers - .get_or_insert_with(key.clone(), || worker_tx.clone()) - .await; - - let backend = Arc::new(MemoryKvBackend::default()); - let table_route_cache = Arc::new(new_table_route_cache( - "pending-rows-flush-deadline-routes".to_string(), - CacheBuilder::new(1).build(), - backend.clone(), - )); - let partition_info_cache = Arc::new(new_partition_info_cache( - "pending-rows-flush-deadline-partitions".to_string(), - CacheBuilder::new(1).build(), - table_route_cache.clone(), - )); - let partition_manager = Arc::new(PartitionRuleManager::new( - backend, - table_route_cache, - partition_info_cache, - )); - let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager { - datanodes: Arc::new(HashMap::new()), - }); - let catalog_manager = MemoryCatalogManager::with_default_setup(); - let (flow_notification_tx, _flow_notification_rx) = Notifier::try_new(1).unwrap(); - let (shutdown, _) = broadcast::channel(1); - - let flush_limiter = FlushLimiter::try_new(1).unwrap(); - start_worker( - key.clone(), - worker_tx.clone(), - workers.clone(), - worker_rx, - shutdown.clone(), - partition_manager, - node_manager, - catalog_manager, - flow_notification_tx, - worker_idle_timeout, - TimingFlushPolicy::try_new(flush_interval, 3).unwrap(), - flush_limiter.clone(), - ); - - // Start the worker, then size-flush a batch halfway to the first - // worker-aligned interval boundary. This arms the reusable timer and - // drains the batch before that deadline is reached. - tokio::task::yield_now().await; - advance(flush_interval / 2).await; - let size_flush_rx = submit_mock_worker_batch(&worker_tx, 3, 1000).await; - let size_flush_result = - receive_mock_flush_result(size_flush_rx, "row threshold did not flush the first batch") - .await; - assert_missing_physical_table(size_flush_result); - - // Submit a low-volume batch before the first batch's old timer would - // expire. It must receive a fresh full interval. - advance(flush_interval / 5).await; - let mut timed_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 2000).await; - - let first_submission = tokio::time::Instant::now(); - advance(flush_interval / 2).await; - let later_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 3000).await; - advance(flush_interval / 2 - Duration::from_millis(1)).await; - for _ in 0..10 { - tokio::task::yield_now().await; - } - assert!(matches!( - timed_flush_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - )); - - advance(Duration::from_millis(1)).await; - let timed_flush_result = receive_mock_flush_result( - timed_flush_rx, - "batch was not flushed one interval after its creation", - ) - .await; - assert_missing_physical_table(timed_flush_result); - - assert_eq!( - first_submission + flush_interval, - tokio::time::Instant::now() - ); - assert_missing_physical_table( - receive_mock_flush_result( - later_flush_rx, - "later submission was not included in the timed flush", - ) - .await, - ); - - // Shutdown retains the existing inline path even with no flush permits. - let _held_permit = flush_limiter.acquire().await.unwrap(); - let shutdown_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 4000).await; - let shutdown_at = tokio::time::Instant::now(); - let _ = shutdown.send(()); - assert_missing_physical_table( - receive_mock_flush_result( - shutdown_flush_rx, - "shutdown incorrectly waited for a flush permit", - ) - .await, - ); - assert_eq!(shutdown_at, tokio::time::Instant::now()); - for _ in 0..10 { - if workers.is_empty().await { - break; - } - tokio::task::yield_now().await; - } - assert!( - workers.is_empty().await, - "worker did not exit after shutdown" - ); - } - - #[tokio::test] - async fn test_flush_batch_physical_uses_mockable_trait_dependencies() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .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)); - } - - #[derive(Default)] - struct AffectedRowsFlushNodeRequester { - 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)) - } - } - - #[tokio::test] - async fn test_flush_batch_physical_returns_actual_affected_rows() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let ctx = session::context::QueryContext::arc(); - - let affected_rows = flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &AffectedRowsFlushNodeRequester { affected_rows: 7 }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .unwrap(); - - assert_eq!(7, affected_rows); - } - - #[tokio::test] - async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - let err = flush_batch_physical( - &table_batches, - "missing_phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { table: None }, - ) - .await - .unwrap_err(); - - assert!( - err.to_string() - .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)); - } - - #[tokio::test] - async fn test_flush_batch_physical_aborts_immediately_on_transform_error() { - let table_batches = vec![ - TableBatch { - table_name: "broken".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("unknown_tag", "host-1", 1000, 1.0)], - row_count: 1, - }, - TableBatch { - table_name: "healthy".to_string(), - table_id: 12, - batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], - row_count: 1, - }, - ]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - let err = flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .unwrap_err(); - - 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)); - } - #[test] fn test_batch_key_groups_by_skip_wal() { use std::collections::HashMap; diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs new file mode 100644 index 00000000000..af9094ef724 --- /dev/null +++ b/src/servers/src/batcher/logical_table/batch.rs @@ -0,0 +1,744 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Instant; + +use catalog::CatalogManagerRef; +use common_batcher::flush_limiter::FlushLimiter; +use common_batcher::notifier::Notifier; +use common_meta::node_manager::NodeManagerRef; +use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_telemetry::{debug, warn}; +use partition::manager::PartitionRuleManagerRef; +use session::context::QueryContextRef; +use snafu::OptionExt; + +use crate::batcher::logical_table::batch_convert::{ + TableBatch, concat_modified_batches, transform_logical_batches_to_physical, +}; +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, +}; +use crate::batcher::logical_table::{ + FlowNotification, PHYSICAL_TABLE_KEY, enqueue_flow_notifications, +}; +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, +}; + +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, +} + +pub(in crate::batcher::logical_table) async fn spawn_flush( + flush: Batch, + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + catalog_manager: CatalogManagerRef, + flow_notification_tx: Notifier, + flush_limiter: FlushLimiter, +) { + match flush_limiter.acquire().await { + Ok(permit) => { + tokio::spawn(async move { + let _permit = permit; + flush_batch_with_managers( + flush, + partition_manager, + node_manager, + catalog_manager, + flow_notification_tx, + ) + .await; + }); + } + Err(err) => { + warn!(err; "Flush semaphore closed, flushing inline"); + flush_batch_with_managers( + flush, + partition_manager, + node_manager, + catalog_manager, + flow_notification_tx, + ) + .await; + } + } +} + +pub(in crate::batcher::logical_table) async fn flush_batch_with_managers( + flush: Batch, + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + catalog_manager: CatalogManagerRef, + flow_notification_tx: Notifier, +) { + let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; + let node_requester = NodeManagerPhysicalFlushAdapter { + node_manager: node_manager.clone(), + }; + let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager }; + flush_batch( + flush, + &partition_provider, + &node_requester, + &catalog_provider, + flow_notification_tx, + ) + .await; +} + +pub(in crate::batcher::logical_table) async fn flush_batch( + flush: Batch, + partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), + node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), + catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), + flow_notification_tx: Notifier, +) { + let Batch { + table_batches, + total_row_count, + db_string, + ctx, + waiters, + } = flush; + let start = Instant::now(); + + // 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) + .to_string(); + let result = flush_batch_physical( + &table_batches, + &physical_table_name, + &ctx, + partition_manager, + node_manager, + catalog_manager, + ) + .await; + + let elapsed = start.elapsed().as_secs_f64(); + FLUSH_ELAPSED.observe(elapsed); + + debug!( + "Pending rows batch flushed, total rows: {}, elapsed time: {}s", + total_row_count, elapsed + ); + + match result { + Ok(affected_rows) => { + FLUSH_TOTAL.inc(); + 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); + } + Err(err) => { + FLUSH_FAILURES.inc(); + FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64); + notify_waiters(waiters, Err(err)); + } + } +} + +/// Flushes a batch of logical table rows by transforming them into the physical table format +/// and writing them to the appropriate datanode regions. +/// +/// This function performs the end-to-end physical flush pipeline: +/// 1. Resolves the physical table metadata and column ID mapping. +/// 2. Fetches the physical table's partition rule. +/// 3. Transforms each logical table batch into the physical (sparse primary key) format. +/// 4. Concatenates all transformed batches into a single combined batch. +/// 5. Splits the combined batch by partition rule and sends region write requests +/// concurrently to the target datanodes. +pub async fn flush_batch_physical( + table_batches: &[TableBatch], + physical_table_name: &str, + ctx: &QueryContextRef, + partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), + node_manager: &(impl PhysicalFlushNodeRequester + ?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 + .with_label_values(&["flush_physical_resolve_table"]) + .start_timer(); + catalog_manager + .physical_table( + ctx.current_catalog(), + &ctx.current_schema(), + physical_table_name, + ctx.as_ref(), + ) + .await? + .with_context(|| error::InternalSnafu { + err_msg: format!( + "Physical table '{}' not found during pending flush", + physical_table_name + ), + })? + }; + + let physical_table_info = physical_table.table_info; + let name_to_ids = physical_table + .col_name_to_ids + .with_context(|| error::InternalSnafu { + err_msg: format!( + "Physical table '{}' has no column IDs for pending flush", + physical_table_name + ), + })?; + + // 2. Get the physical table's partition rule (one lookup instead of N) + let partition_rule = { + let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_physical_fetch_partition_rule"]) + .start_timer(); + partition_manager + .find_table_partition_rule(physical_table_info.as_ref()) + .await? + }; + let partition_columns = partition_rule.partition_columns(); + let partition_columns_set: HashSet<&str> = + 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)?; + + // 4. Concatenate all modified batches (all share the same physical schema) + let combined_batch = 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( + 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 +} + +pub(in crate::batcher::logical_table) fn notify_waiters( + waiters: Vec, + result: Result<()>, +) { + let shared_result = result.map_err(Arc::new); + for waiter in waiters { + let _ = waiter.response_tx.send(match &shared_result { + Ok(()) => Ok(()), + Err(error) => Err(Arc::clone(error)), + }); + // waiter._permit is dropped here, releasing the inflight semaphore slot + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::future::poll_fn; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + 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; + use async_trait::async_trait; + use catalog::error::Result as CatalogResult; + use common_batcher::notifier::Notifier; + use common_batcher::request_limiter::RequestLimiter; + use common_meta::cache::TableFlownodeSetCacheRef; + use common_meta::node_manager::NodeManagerRef; + 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}; + + use crate::batcher::logical_table::batch::{ + Batch, flush_batch, flush_batch_physical, notify_waiters, + }; + 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, + PhysicalTableMetadata, + }; + use crate::batcher::logical_table::start_flow_notification_worker; + use crate::batcher::logical_table::test_util::{ + FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch, + mock_table_flownode_cache, + }; + use crate::error; + use crate::error::Error; + + fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata { + let schema = Arc::new( + DtSchema::try_new(vec![ + DtColumnSchema::new( + "__primary_key", + datatypes::prelude::ConcreteDataType::binary_datatype(), + false, + ), + DtColumnSchema::new( + "greptime_timestamp", + datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + DtColumnSchema::new( + "greptime_value", + datatypes::prelude::ConcreteDataType::float64_datatype(), + true, + ), + DtColumnSchema::new( + "tag1", + datatypes::prelude::ConcreteDataType::string_datatype(), + true, + ), + ]) + .unwrap(), + ); + let mut table_info = test_table_info(table_id, "phy", "public", "greptime", schema); + table_info.meta.column_ids = vec![0, 1, 2, 3]; + + PhysicalTableMetadata { + table_info: Arc::new(table_info), + col_name_to_ids: Some(HashMap::from([("tag1".to_string(), 3)])), + } + } + + struct MockFlushCatalogProvider { + table: Option, + } + + #[async_trait] + impl PhysicalFlushCatalogProvider for MockFlushCatalogProvider { + async fn physical_table( + &self, + _catalog: &str, + _schema: &str, + _table_name: &str, + _query_ctx: &session::context::QueryContext, + ) -> CatalogResult> { + Ok(self.table.clone()) + } + } + + struct SingleRegionPartitionRule; + + impl PartitionRule for SingleRegionPartitionRule { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn partition_columns(&self) -> &[String] { + &[] + } + + fn find_region( + &self, + _values: &[datatypes::prelude::Value], + ) -> partition::error::Result { + unimplemented!() + } + + fn split_record_batch( + &self, + record_batch: &RecordBatch, + ) -> partition::error::Result> + { + Ok(HashMap::from([( + 1, + RegionMask::new( + arrow::array::BooleanArray::from(vec![true; record_batch.num_rows()]), + record_batch.num_rows(), + ), + )])) + } + } + + struct MockFlushPartitionProvider { + partition_rule_calls: Arc, + region_leader_calls: Arc, + } + + #[async_trait] + impl PhysicalFlushPartitionProvider for MockFlushPartitionProvider { + async fn find_table_partition_rule( + &self, + _table_info: &table::metadata::TableInfo, + ) -> PartitionResult { + 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(), + }) + } + } + + #[derive(Default)] + struct MockFlushNodeRequester { + writes: Arc, + fail: bool, + } + + #[async_trait] + impl PhysicalFlushNodeRequester for MockFlushNodeRequester { + async fn handle( + &self, + _peer: &Peer, + _request: RegionRequest, + ) -> 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)) + } + } + + fn mock_flow_notification_sender( + cache: TableFlownodeSetCacheRef, + node_manager: NodeManagerRef, + ) -> Notifier { + let (tx, rx) = Notifier::try_new(16).unwrap(); + start_flow_notification_worker(rx, cache, node_manager); + tx + } + + #[tokio::test] + async fn test_flush_batch_notifies_flownode_after_successful_physical_write() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, peer).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); + let ctx = session::context::QueryContext::arc(); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let writes = Arc::new(AtomicUsize::new(0)); + + flush_batch( + Batch { + table_batches, + total_row_count: 1, + db_string: ctx.get_db_string(), + ctx, + waiters: Vec::new(), + }, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &MockFlushNodeRequester { + writes: writes.clone(), + fail: false, + }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + flow_notification_tx, + ) + .await; + + assert_eq!(1, writes.load(Ordering::SeqCst)); + let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + vec![api::v1::flow::DirtyWindowRequest { + table_id, + timestamps: vec![1000], + time_ranges: Vec::new(), + }], + requests.requests + ); + } + + #[tokio::test] + async fn test_flush_batch_does_not_notify_flownode_after_physical_write_error() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, peer).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); + let ctx = session::context::QueryContext::arc(); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let writes = Arc::new(AtomicUsize::new(0)); + + flush_batch( + Batch { + table_batches, + total_row_count: 1, + db_string: ctx.get_db_string(), + ctx, + waiters: Vec::new(), + }, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &MockFlushNodeRequester { + writes: writes.clone(), + fail: true, + }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + flow_notification_tx, + ) + .await; + + assert_eq!(1, writes.load(Ordering::SeqCst)); + assert!( + tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_cancelled_waiter_retains_request_slot_until_notification() { + let limiter = RequestLimiter::try_new(1).unwrap(); + let (response_tx, response_rx) = oneshot::channel(); + let waiter = FlushWaiter { + response_tx, + _permit: limiter.acquire().await.unwrap(), + }; + drop(response_rx); + let next = limiter.acquire(); + tokio::pin!(next); + assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await); + notify_waiters(vec![waiter], Ok(())); + let _permit = next.await.unwrap(); + } + + #[tokio::test] + async fn test_flush_batch_physical_uses_mockable_trait_dependencies() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .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)); + } + + #[derive(Default)] + struct AffectedRowsFlushNodeRequester { + 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)) + } + } + + #[tokio::test] + async fn test_flush_batch_physical_returns_actual_affected_rows() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let ctx = session::context::QueryContext::arc(); + + let affected_rows = flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &AffectedRowsFlushNodeRequester { affected_rows: 7 }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .unwrap(); + + assert_eq!(7, affected_rows); + } + + #[tokio::test] + async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + let err = flush_batch_physical( + &table_batches, + "missing_phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { table: None }, + ) + .await + .unwrap_err(); + + assert!( + err.to_string() + .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)); + } + + #[tokio::test] + async fn test_flush_batch_physical_aborts_immediately_on_transform_error() { + let table_batches = vec![ + TableBatch { + table_name: "broken".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("unknown_tag", "host-1", 1000, 1.0)], + row_count: 1, + }, + TableBatch { + table_name: "healthy".to_string(), + table_id: 12, + batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], + row_count: 1, + }, + ]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + let err = flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .unwrap_err(); + + 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)); + } +} diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs new file mode 100644 index 00000000000..e7bd33c1975 --- /dev/null +++ b/src/servers/src/batcher/logical_table/pending_worker.rs @@ -0,0 +1,656 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use catalog::CatalogManagerRef; +use common_batcher::flush_limiter::FlushLimiter; +use common_batcher::flush_policy::FlushTrigger; +use common_batcher::flush_policy::timing::TimingFlushPolicy; +use common_batcher::notifier::Notifier; +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 partition::manager::PartitionRuleManagerRef; +use session::context::QueryContextRef; +use table::metadata::TableId; +use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot}; + +use crate::batcher::logical_table::batch::{Batch, flush_batch_with_managers, spawn_flush}; +use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; +use crate::batcher::logical_table::{BatchKey, FlowNotification}; +use crate::error::Error; +use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS}; + +pub(in crate::batcher::logical_table) struct PendingBatch { + pub(in crate::batcher::logical_table) tables: HashMap, + pub(in crate::batcher::logical_table) db_string: String, + pub(in crate::batcher::logical_table) ctx: QueryContextRef, +} + +pub(in crate::batcher::logical_table) struct FlushWaiter { + pub(in crate::batcher::logical_table) response_tx: + oneshot::Sender>>, + pub(in crate::batcher::logical_table) _permit: Arc, +} + +#[derive(Clone)] +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 { + Submit { + table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, + total_rows: usize, + ctx: QueryContextRef, + response_tx: oneshot::Sender>>, + _permit: Arc, + }, + #[cfg(test)] + Ack { ack_tx: oneshot::Sender<()> }, +} + +impl PendingBatch { + pub(in crate::batcher::logical_table) fn new(ctx: QueryContextRef) -> Self { + let db_string = ctx.get_db_string(); + Self { + tables: HashMap::new(), + db_string, + ctx, + } + } + + pub(in crate::batcher::logical_table) fn add_table_batch( + &mut self, + table_name: String, + table_id: TableId, + record_batch: RecordBatchWithTsIdx, + ) { + let entry = self.tables.entry(table_id).or_insert_with(|| TableBatch { + table_name, + table_id, + batches: Vec::new(), + row_count: 0, + }); + entry.row_count += record_batch.batch.num_rows(); + entry.batches.push(record_batch); + } +} + +#[allow(clippy::too_many_arguments)] +pub(in crate::batcher::logical_table) fn start_worker( + key: BatchKey, + worker_tx: mpsc::Sender, + workers: Arc>, + mut rx: mpsc::Receiver, + shutdown: broadcast::Sender<()>, + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + catalog_manager: CatalogManagerRef, + flow_notification_tx: Notifier, + worker_idle_timeout: Duration, + flush_policy: TimingFlushPolicy, + flush_limiter: FlushLimiter, +) { + tokio::spawn(async move { + // The business batch and pending flush are populated and drained together. + let mut batch = None; + let mut pending_flush = PendingCore::new(flush_policy); + let mut shutdown_rx = shutdown.subscribe(); + let idle_deadline = tokio::time::Instant::now() + worker_idle_timeout; + let idle_timer = tokio::time::sleep_until(idle_deadline); + tokio::pin!(idle_timer); + + loop { + tokio::select! { + cmd = rx.recv() => { + match cmd { + Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit }) => { + let submitted_at = tokio::time::Instant::now(); + idle_timer.as_mut().reset(submitted_at + worker_idle_timeout); + + pending_flush.submit( + FlushWaiter { response_tx, _permit }, + total_rows, + ); + let pending_batch = batch.get_or_insert_with(||{ + PENDING_BATCHES.inc(); + PendingBatch::new(ctx) + }); + + for (table_name, table_id, record_batch) in table_batches { + pending_batch.add_table_batch(table_name, table_id, record_batch); + } + + PENDING_ROWS.add(total_rows as i64); + + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) { + spawn_flush( + flush, + partition_manager.clone(), + node_manager.clone(), + catalog_manager.clone(), + flow_notification_tx.clone(), + flush_limiter.clone(), + ).await; + } + } + None => { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { + flush_batch_with_managers( + flush, + partition_manager.clone(), + node_manager.clone(), + catalog_manager.clone(), + flow_notification_tx.clone(), + ).await; + } + break; + } + #[cfg(test)] + Some(WorkerCommand::Ack { ack_tx }) => { + let _ = ack_tx.send(()); + } + } + } + _ = &mut idle_timer => { + if !should_close_worker_on_idle_timeout( + pending_flush.total_rows(), + rx.len(), + ) { + idle_timer + .as_mut() + .reset(tokio::time::Instant::now() + worker_idle_timeout); + continue; + } + + debug!( + "Closing idle pending rows worker due to timeout: catalog={}, schema={}, physical_table={}", + key.catalog, + key.schema, + key.physical_table + ); + break; + } + _ = pending_flush.wait_flush() => { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) { + spawn_flush( + flush, + partition_manager.clone(), + node_manager.clone(), + catalog_manager.clone(), + flow_notification_tx.clone(), + flush_limiter.clone(), + ).await; + } + } + _ = shutdown_rx.recv() => { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { + flush_batch_with_managers( + flush, + partition_manager.clone(), + node_manager.clone(), + catalog_manager.clone(), + flow_notification_tx.clone(), + ).await; + } + break; + } + } + } + + remove_worker_if_same_channel(workers.as_ref(), &key, &worker_tx).await; + }); +} + +pub(in crate::batcher::logical_table) async fn remove_worker_if_same_channel( + workers: &WorkerRegistry, + key: &BatchKey, + worker_tx: &mpsc::Sender, +) -> bool { + if workers.remove_if_same(key, worker_tx).await { + PENDING_WORKERS.set(workers.len().await as i64); + true + } else { + false + } +} + +pub(in crate::batcher::logical_table) fn should_close_worker_on_idle_timeout( + total_row_count: usize, + queued_requests: usize, +) -> bool { + total_row_count == 0 && queued_requests == 0 +} + +/// Transfers the ready batch, or drains unconditionally when no trigger is given. +/// Execution and flush permits remain owned by the caller. +pub(in crate::batcher::logical_table) fn drain_batch( + batch: &mut Option, + pending_flush: &mut PendingCore, + trigger: Option, +) -> Option { + let total_row_count = pending_flush.total_rows(); + let waiters = match trigger { + Some(trigger) => pending_flush.take_ready(trigger)?, + None => pending_flush.take_pending()?, + }; + let batch = batch.take()?; + + if total_row_count == 0 { + return None; + } + + let table_batches = batch.tables.into_values().collect(); + + PENDING_ROWS.sub(total_row_count as i64); + PENDING_BATCHES.dec(); + + Some(Batch { + table_batches, + total_row_count, + db_string: batch.db_string, + ctx: batch.ctx, + waiters, + }) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::time::Duration; + + use catalog::memory::MemoryCatalogManager; + use common_batcher::flush_limiter::FlushLimiter; + use common_batcher::flush_policy::FlushTrigger; + use common_batcher::flush_policy::timing::TimingFlushPolicy; + use common_batcher::notifier::Notifier; + use common_batcher::pending_worker::PendingWorker as PendingCore; + use common_batcher::worker_registry::WorkerRegistry; + use common_meta::cache::new_table_route_cache; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::node_manager::NodeManagerRef; + use moka::future::CacheBuilder; + use partition::cache::new_partition_info_cache; + use partition::manager::PartitionRuleManager; + use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; + use tokio::time::advance; + + use crate::batcher::logical_table::BatchKey; + use crate::batcher::logical_table::batch_convert::TableBatch; + use crate::batcher::logical_table::pending_worker::{ + FlushWaiter, PendingBatch, WorkerCommand, drain_batch, remove_worker_if_same_channel, + should_close_worker_on_idle_timeout, start_worker, + }; + use crate::batcher::logical_table::test_util::{ + ConcurrentMockNodeManager, mock_aligned_tag_batch, + }; + use crate::error::Error; + + #[tokio::test] + async fn test_drain_batch_takes_initialized_pending_batch_from_option() { + let ctx = session::context::QueryContext::arc(); + let (response_tx, _response_rx) = oneshot::channel(); + let permit = Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap(); + let mut pending_flush = + PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 1).unwrap()); + pending_flush.submit( + FlushWaiter { + response_tx, + _permit: Arc::new(permit), + }, + 1, + ); + let mut batch = Some(PendingBatch { + tables: HashMap::from([( + 42, + TableBatch { + table_name: "cpu".to_string(), + table_id: 42, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }, + )]), + db_string: ctx.get_db_string(), + ctx: ctx.clone(), + }); + + let flush = drain_batch( + &mut batch, + &mut pending_flush, + Some(FlushTrigger::Submission), + ) + .unwrap(); + + assert!(batch.is_none()); + assert!(pending_flush.is_empty()); + assert_eq!(0, pending_flush.total_rows()); + assert_eq!(1, flush.waiters.len()); + assert_eq!(1, flush.total_row_count); + assert_eq!(1, flush.table_batches.len()); + assert_eq!(ctx.get_db_string(), flush.db_string); + assert_eq!(ctx.current_catalog(), flush.ctx.current_catalog()); + } + + #[tokio::test] + async fn test_drain_batch_preserves_unready_state_and_clears_zero_rows() { + for total_rows in [0, 1] { + let mut pending_flush = + PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 2).unwrap()); + let mut batch = Some(PendingBatch::new(session::context::QueryContext::arc())); + let semaphore = Arc::new(Semaphore::new(1)); + let (response_tx, mut response_rx) = oneshot::channel(); + pending_flush.submit( + FlushWaiter { + response_tx, + _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()), + }, + total_rows, + ); + assert!( + drain_batch( + &mut batch, + &mut pending_flush, + Some(FlushTrigger::Submission) + ) + .is_none() + ); + assert!(batch.is_some()); + assert!(!pending_flush.is_empty()); + assert_eq!(total_rows, pending_flush.total_rows()); + assert_eq!(0, semaphore.available_permits()); + + let drained = drain_batch(&mut batch, &mut pending_flush, None); + assert!(batch.is_none()); + assert!(pending_flush.is_empty()); + assert_eq!(0, pending_flush.total_rows()); + if total_rows == 0 { + assert!(drained.is_none()); + assert_eq!(1, semaphore.available_permits()); + assert!(matches!( + response_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + } else { + let drained = drained.unwrap(); + assert_eq!(1, drained.total_row_count); + assert_eq!(1, drained.waiters.len()); + assert_eq!(0, semaphore.available_permits()); + drop(drained); + assert_eq!(1, semaphore.available_permits()); + } + } + } + + #[test] + fn test_pending_batch_keeps_same_name_batches_with_distinct_table_ids() { + let ctx = session::context::QueryContext::arc(); + let mut pending_batch = PendingBatch::new(ctx); + + pending_batch.add_table_batch( + "cpu".to_string(), + 42, + mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0), + ); + pending_batch.add_table_batch( + "cpu".to_string(), + 43, + mock_aligned_tag_batch("tag1", "host-1", 2000, 2.0), + ); + + assert_eq!(2, pending_batch.tables.len()); + assert_eq!(42, pending_batch.tables[&42].table_id); + assert_eq!(43, pending_batch.tables[&43].table_id); + assert_eq!("cpu", pending_batch.tables[&42].table_name); + assert_eq!("cpu", pending_batch.tables[&43].table_name); + } + + #[tokio::test] + async fn test_remove_worker_if_same_channel_removes_matching_entry() { + let workers = WorkerRegistry::new(); + let key = BatchKey { + catalog: "greptime".to_string(), + schema: "public".to_string(), + physical_table: "phy".to_string(), + skip_wal: false, + }; + + let (tx, _rx) = mpsc::channel::(1); + workers.get_or_insert_with(key.clone(), || tx.clone()).await; + + assert!(remove_worker_if_same_channel(&workers, &key, &tx).await); + assert!(workers.is_empty().await); + } + + #[tokio::test] + async fn test_remove_worker_if_same_channel_keeps_newer_entry() { + let workers = WorkerRegistry::new(); + let key = BatchKey { + catalog: "greptime".to_string(), + schema: "public".to_string(), + physical_table: "phy".to_string(), + skip_wal: false, + }; + + let (stale_tx, _stale_rx) = mpsc::channel::(1); + let (fresh_tx, _fresh_rx) = mpsc::channel::(1); + workers + .get_or_insert_with(key.clone(), || fresh_tx.clone()) + .await; + + assert!(!remove_worker_if_same_channel(&workers, &key, &stale_tx).await); + assert!(workers.get(&key).await.is_some()); + assert!(workers.get(&key).await.unwrap().same_channel(&fresh_tx)); + } + + #[test] + fn test_worker_idle_timeout_close_decision() { + assert!(should_close_worker_on_idle_timeout(0, 0)); + assert!(!should_close_worker_on_idle_timeout(1, 0)); + assert!(!should_close_worker_on_idle_timeout(0, 1)); + } + + const WORKER_TEST_TIMEOUT: Duration = Duration::from_secs(30); + + async fn submit_mock_worker_batch( + worker_tx: &mpsc::Sender, + total_rows: usize, + timestamp: i64, + ) -> oneshot::Receiver>> { + let (response_tx, response_rx) = oneshot::channel(); + let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); + worker_tx + .send(WorkerCommand::Submit { + table_batches: vec![( + "cpu".to_string(), + 42, + mock_aligned_tag_batch("tag1", "host-1", timestamp, 1.0), + )], + total_rows, + ctx: session::context::QueryContext::arc(), + response_tx, + _permit: Arc::new(permit), + }) + .await + .unwrap(); + + // The channel is FIFO, so the ack proves the worker has dequeued and + // processed the submission (anchoring the flush deadline) before the + // caller advances virtual time. + let (ack_tx, ack_rx) = oneshot::channel(); + worker_tx.send(WorkerCommand::Ack { ack_tx }).await.unwrap(); + ack_rx + .await + .expect("worker exited before acking the submitted batch"); + + response_rx + } + + async fn receive_mock_flush_result( + response_rx: oneshot::Receiver>>, + context: &str, + ) -> std::result::Result<(), Arc> { + // Under paused time the timeout auto-advances the clock and fires + // deterministically if the flush never completes. + tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx) + .await + .unwrap_or_else(|_| panic!("{context}")) + .expect("flush result channel closed without a result") + } + + fn assert_missing_physical_table(result: std::result::Result<(), Arc>) { + let err = result.expect_err("the empty catalog should make the flush fail"); + assert!( + matches!( + err.as_ref(), + Error::Internal { err_msg } + if err_msg.contains("not found during pending flush") + ), + "unexpected flush error: {err}" + ); + } + + #[tokio::test(start_paused = true)] + async fn test_worker_preserves_first_deadline_and_inline_shutdown() { + let flush_interval = Duration::from_secs(10); + let worker_idle_timeout = Duration::from_secs(30); + let key = BatchKey { + catalog: "greptime".to_string(), + schema: "public".to_string(), + physical_table: "phy".to_string(), + skip_wal: false, + }; + let workers = Arc::new(WorkerRegistry::new()); + let (worker_tx, worker_rx) = mpsc::channel(1); + workers + .get_or_insert_with(key.clone(), || worker_tx.clone()) + .await; + + let backend = Arc::new(MemoryKvBackend::default()); + let table_route_cache = Arc::new(new_table_route_cache( + "pending-rows-flush-deadline-routes".to_string(), + CacheBuilder::new(1).build(), + backend.clone(), + )); + let partition_info_cache = Arc::new(new_partition_info_cache( + "pending-rows-flush-deadline-partitions".to_string(), + CacheBuilder::new(1).build(), + table_route_cache.clone(), + )); + let partition_manager = Arc::new(PartitionRuleManager::new( + backend, + table_route_cache, + partition_info_cache, + )); + let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager { + datanodes: Arc::new(HashMap::new()), + }); + let catalog_manager = MemoryCatalogManager::with_default_setup(); + let (flow_notification_tx, _flow_notification_rx) = Notifier::try_new(1).unwrap(); + let (shutdown, _) = broadcast::channel(1); + + let flush_limiter = FlushLimiter::try_new(1).unwrap(); + start_worker( + key.clone(), + worker_tx.clone(), + workers.clone(), + worker_rx, + shutdown.clone(), + partition_manager, + node_manager, + catalog_manager, + flow_notification_tx, + worker_idle_timeout, + TimingFlushPolicy::try_new(flush_interval, 3).unwrap(), + flush_limiter.clone(), + ); + + // Start the worker, then size-flush a batch halfway to the first + // worker-aligned interval boundary. This arms the reusable timer and + // drains the batch before that deadline is reached. + tokio::task::yield_now().await; + advance(flush_interval / 2).await; + let size_flush_rx = submit_mock_worker_batch(&worker_tx, 3, 1000).await; + let size_flush_result = + receive_mock_flush_result(size_flush_rx, "row threshold did not flush the first batch") + .await; + assert_missing_physical_table(size_flush_result); + + // Submit a low-volume batch before the first batch's old timer would + // expire. It must receive a fresh full interval. + advance(flush_interval / 5).await; + let mut timed_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 2000).await; + + let first_submission = tokio::time::Instant::now(); + advance(flush_interval / 2).await; + let later_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 3000).await; + advance(flush_interval / 2 - Duration::from_millis(1)).await; + for _ in 0..10 { + tokio::task::yield_now().await; + } + assert!(matches!( + timed_flush_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + + advance(Duration::from_millis(1)).await; + let timed_flush_result = receive_mock_flush_result( + timed_flush_rx, + "batch was not flushed one interval after its creation", + ) + .await; + assert_missing_physical_table(timed_flush_result); + + assert_eq!( + first_submission + flush_interval, + tokio::time::Instant::now() + ); + assert_missing_physical_table( + receive_mock_flush_result( + later_flush_rx, + "later submission was not included in the timed flush", + ) + .await, + ); + + // Shutdown retains the existing inline path even with no flush permits. + let _held_permit = flush_limiter.acquire().await.unwrap(); + let shutdown_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 4000).await; + let shutdown_at = tokio::time::Instant::now(); + let _ = shutdown.send(()); + assert_missing_physical_table( + receive_mock_flush_result( + shutdown_flush_rx, + "shutdown incorrectly waited for a flush permit", + ) + .await, + ); + assert_eq!(shutdown_at, tokio::time::Instant::now()); + for _ in 0..10 { + if workers.is_empty().await { + break; + } + tokio::task::yield_now().await; + } + assert!( + workers.is_empty().await, + "worker did not exit after shutdown" + ); + } +} From 263e229103d955d34bb977c7301d324933d891fc Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 21 Sep 2026 09:10:01 +0000 Subject: [PATCH 41/63] feat: add experimental Metric export to V2 snapshots (#9233) * feat: add experimental Metric export to V2 snapshots Signed-off-by: jeremyhi * fix: validate the complete Metric export capability response Signed-off-by: jeremyhi * test: construct portable file URLs for Metric export fixtures Signed-off-by: jeremyhi * refactor: address Metric export review nits Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 1 + config/config.md | 2 + config/frontend.example.toml | 4 + config/standalone.example.toml | 4 + src/cli/src/data/export_v2.rs | 9 + src/cli/src/data/export_v2/command.rs | 59 ++- src/cli/src/data/export_v2/coordinator.rs | 183 ++++++- src/cli/src/data/export_v2/data.rs | 4 + src/cli/src/data/export_v2/error.rs | 11 +- src/cli/src/data/snapshot_storage.rs | 206 ++++++++ src/cli/src/database.rs | 18 +- src/frontend/src/frontend.rs | 6 + src/frontend/src/instance.rs | 126 +++++ src/frontend/src/instance/builder.rs | 1 + src/frontend/src/instance/export_database.rs | 68 +++ src/operator/src/statement.rs | 2 +- src/standalone/src/options.rs | 7 + tests-integration/Cargo.toml | 1 + tests-integration/src/standalone.rs | 10 + .../tests/export_logical_tables.rs | 484 ++++++++++++++++++ tests-integration/tests/http.rs | 1 + 21 files changed, 1172 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 36ee6d77ffd..dc0dcf2561c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15125,6 +15125,7 @@ dependencies = [ "catalog", "chrono", "clap", + "cli", "client", "cmd", "common-base", diff --git a/config/config.md b/config/config.md index 651cbec6e85..7385c8efb0d 100644 --- a/config/config.md +++ b/config/config.md @@ -15,6 +15,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | +| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | @@ -278,6 +279,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | +| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 95550931d46..043da57db74 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -11,6 +11,10 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true +## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. +## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. +#+ experimental_metric_export = false + ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 57893881146..cc4be3f5930 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -11,6 +11,10 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true +## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. +## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. +#+ experimental_metric_export = false + ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index 9bdf3c63f5d..f9a3c2be095 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -37,6 +37,15 @@ //! --start-time 2025-01-01T00:00:00Z \ //! --end-time 2025-01-31T23:59:59Z //! ``` +//! +//! `--experimental-metric-export` enables shared Metric physical scans for Parquet. +//! Enable `experimental_metric_export = true` on the frontend or standalone server +//! and use an endpoint whose frontends all support and enable this option. +//! The command checks the server capability before modifying the snapshot. +//! Each snapshot path belongs to one export task. Resume requires stable source +//! schemas/data and confirmation that the previous export and storage writes ended; +//! an HTTP timeout does not establish that. Only unfinished chunks are cleaned and +//! rerun. Completed chunks and the V2 snapshot/import format remain unchanged. mod chunker; mod command; diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index 4df741db347..fc958e9e3e0 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -20,9 +20,11 @@ use std::time::Duration; use async_trait::async_trait; use clap::{Parser, Subcommand}; +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_error::ext::BoxedError; use common_telemetry::info; use serde_json::Value; +use servers::http::{ColumnSchema, GreptimeQueryOutput, OutputSchema}; use snafu::{OptionExt, ResultExt}; use crate::Tool; @@ -289,6 +291,11 @@ pub struct ExportCreateCommand { #[clap(long, value_enum, default_value = "parquet")] format: DataFormat, + /// Use shared Metric physical scans (Parquet only). Resume requires that the previous + /// export and storage writes have ended; an HTTP timeout does not establish this. + #[clap(long)] + experimental_metric_export: bool, + /// Delete existing snapshot and recreate. #[clap(long)] force: bool, @@ -379,9 +386,6 @@ impl ExportCreateCommand { Some(self.schemas.clone()) }; - // Build storage - let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; - // Build database client let proxy = parse_proxy_opts(self.proxy.clone(), self.no_proxy)?; let database_client = DatabaseClient::new( @@ -393,12 +397,43 @@ impl ExportCreateCommand { self.no_proxy, ); + // The filesystem storage constructor can create the snapshot root. + if self.experimental_metric_export { + if self.format != DataFormat::Parquet { + return crate::data::export_v2::error::MetricExportUnavailableSnafu + .fail() + .map_err(BoxedError::new); + } + let capability = database_client + .sql_response( + "SHOW VARIABLES experimental_metric_export", + DEFAULT_SCHEMA_NAME, + ) + .await + .context(DatabaseSnafu) + .map_err(BoxedError::new)?; + let expected_schema = OutputSchema::new(vec![ColumnSchema::new( + "EXPERIMENTAL_METRIC_EXPORT".to_string(), + "String".to_string(), + )]); + if !matches!(capability.output(), [GreptimeQueryOutput::Records(records)] + if records.schema() == &expected_schema + && records.rows() == &vec![vec![Value::String("true".to_string())]]) + { + return crate::data::export_v2::error::MetricExportUnavailableSnafu + .fail() + .map_err(BoxedError::new); + } + } + let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; + Ok(Box::new(ExportCreate { config: ExportConfig { catalog: self.catalog.clone(), schemas, schema_only: self.schema_only, format: self.format, + experimental_metric_export: self.experimental_metric_export, force: self.force, time_range, chunk_time_window: self.chunk_time_window, @@ -426,6 +461,7 @@ struct ExportConfig { schemas: Option>, schema_only: bool, format: DataFormat, + experimental_metric_export: bool, force: bool, time_range: TimeRange, chunk_time_window: Option, @@ -504,6 +540,8 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, + experimental_metric_export: self.config.experimental_metric_export, + resume: true, }, progress.as_ref(), ) @@ -533,6 +571,14 @@ impl ExportCreate { self.config.chunk_time_window, )?; + if self.config.experimental_metric_export { + for chunk in &manifest.chunks { + self.storage + .prepare_export_chunk(&schema_names, chunk.id, false) + .await?; + } + } + // 4. Write schema files self.storage.write_schema(&schema_snapshot).await?; info!("Exported {} schemas", schema_snapshot.schemas.len()); @@ -566,6 +612,8 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, + experimental_metric_export: self.config.experimental_metric_export, + resume: false, }, progress.as_ref(), ) @@ -1744,6 +1792,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1781,6 +1830,7 @@ mod tests { ]), schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1813,6 +1863,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range, chunk_time_window: Some(Duration::from_secs(3600)), @@ -1846,6 +1897,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Csv, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1881,6 +1933,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::new(Some(start), Some(start)), chunk_time_window: None, diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index 64aa4ce9946..b70897337a7 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_telemetry::info; +use common_telemetry::{error, info}; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -39,6 +39,8 @@ struct ExportContext<'a> { schemas: Vec, format: DataFormat, parallelism: usize, + experimental_metric_export: bool, + resume: bool, } pub struct ExportDataOptions<'a> { @@ -46,6 +48,8 @@ pub struct ExportDataOptions<'a> { pub storage_config: &'a ObjectStoreConfig, pub parallelism: usize, pub chunk_parallelism: usize, + pub experimental_metric_export: bool, + pub resume: bool, } pub async fn export_data( @@ -68,6 +72,8 @@ pub async fn export_data( schemas: manifest.schemas.clone(), format: manifest.format, parallelism: options.parallelism, + experimental_metric_export: options.experimental_metric_export, + resume: options.resume, }; // One progress unit per chunk. Already completed/skipped chunks from a @@ -88,11 +94,11 @@ pub async fn export_data( export_data_serial(&context, storage, manifest, progress).await } else { export_data_concurrent( - &context, storage, manifest, options.chunk_parallelism, progress, + |id, range| export_chunk(&context, id, range), ) .await }; @@ -134,11 +140,17 @@ async fn export_data_serial( }; manifest.touch(); - storage.write_manifest(manifest).await?; - // The chunk is finalized (completed, skipped, or failed) and persisted. - progress.inc(1); - - result?; + let persistence = storage.write_manifest(manifest).await; + if persistence.is_ok() { + progress.inc(1); + } + if let Err(err) = result { + if let Err(secondary) = persistence { + error!(secondary; "Failed to persist failed export chunk"); + } + return Err(err); + } + persistence?; } Ok(()) @@ -155,39 +167,44 @@ async fn export_data_serial( /// On the first chunk failure we stop scheduling new chunks but let already /// in-flight chunks finish and persist their final status, then return the /// first error. -async fn export_data_concurrent( - context: &ExportContext<'_>, +async fn export_data_concurrent( storage: &dyn SnapshotStorage, manifest: &mut Manifest, chunk_parallelism: usize, progress: &dyn ProgressReporter, -) -> Result<()> { + export: F, +) -> Result<()> +where + F: Fn(u32, TimeRange) -> Fut, + Fut: std::future::Future>>, +{ let mut pending = FuturesUnordered::new(); let mut next_idx = 0; let mut first_error: Option = None; loop { - let mut scheduled = false; - + let mut admitted = Vec::new(); // Schedule eligible chunks in order up to the parallelism limit. Once a // failure is seen, stop scheduling but keep draining in-flight chunks. - while first_error.is_none() && pending.len() < chunk_parallelism { + while first_error.is_none() && pending.len() + admitted.len() < chunk_parallelism { let Some(idx) = next_eligible_chunk(manifest, &mut next_idx) else { break; }; let (chunk_id, time_range) = mark_chunk_in_progress(manifest, idx); - scheduled = true; - - pending.push(async move { - let result = export_chunk(context, chunk_id, time_range).await; - (idx, result) - }); + admitted.push((idx, chunk_id, time_range)); } - - if scheduled { + if !admitted.is_empty() { manifest.touch(); - storage.write_manifest(manifest).await?; + match storage.write_manifest(manifest).await { + Ok(()) => { + for (idx, chunk_id, time_range) in admitted { + let export = &export; + pending.push(async move { (idx, export(chunk_id, time_range).await) }); + } + } + Err(err) => first_error = Some(err), + } } let Some((idx, export_result)) = pending.next().await else { @@ -204,9 +221,11 @@ async fn export_data_concurrent( } } manifest.touch(); - storage.write_manifest(manifest).await?; - // The chunk is finalized (completed, skipped, or failed) and persisted. - progress.inc(1); + match storage.write_manifest(manifest).await { + Ok(()) => progress.inc(1), + Err(err) if first_error.is_none() => first_error = Some(err), + Err(err) => error!(err; "Failed to persist export chunk while draining"), + } } match first_error { @@ -257,12 +276,19 @@ async fn export_chunk( chunk_id: u32, time_range: TimeRange, ) -> Result> { + if context.experimental_metric_export { + context + .storage + .prepare_export_chunk(&context.schemas, chunk_id, context.resume) + .await?; + } let scheme = StorageScheme::from_uri(context.snapshot_uri)?; let needs_dir = matches!(scheme, StorageScheme::File); let copy_options = CopyOptions { format: context.format, time_range, parallelism: context.parallelism, + experimental_metric_export: context.experimental_metric_export, }; for schema in &context.schemas { @@ -326,6 +352,113 @@ mod tests { manifest } + struct FailingManifest { + fail_at: usize, + writes: std::sync::atomic::AtomicUsize, + release: tokio::sync::Notify, + } + + #[async_trait::async_trait] + impl SnapshotStorage for FailingManifest { + async fn write_manifest(&self, manifest: &Manifest) -> Result<()> { + use std::sync::atomic::Ordering; + + let write = self.writes.fetch_add(1, Ordering::SeqCst) + 1; + if write == 1 { + assert_eq!(manifest.in_progress_count(), 2); + } + if write == self.fail_at { + self.release.notify_one(); + return crate::data::export_v2::error::IoSnafu { + operation: "injected manifest failure", + error: std::io::Error::other("manifest failure"), + } + .fail(); + } + Ok(()) + } + async fn exists(&self) -> Result { + unreachable!() + } + async fn read_manifest(&self) -> Result { + unreachable!() + } + async fn write_schema( + &self, + _: &crate::data::export_v2::schema::SchemaSnapshot, + ) -> Result<()> { + unreachable!() + } + async fn write_text(&self, _: &str, _: &str) -> Result<()> { + unreachable!() + } + async fn read_text(&self, _: &str) -> Result { + unreachable!() + } + async fn create_dir_all(&self, _: &str) -> Result<()> { + unreachable!() + } + async fn list_files_recursive(&self, _: &str) -> Result> { + unreachable!() + } + async fn delete_snapshot(&self) -> Result<()> { + unreachable!() + } + } + + #[tokio::test] + async fn test_manifest_failure_stops_admission_and_drains_started_chunks() { + use std::sync::atomic::{AtomicUsize, Ordering}; + for fail_at in [1, 2, 3] { + let storage = FailingManifest { + fail_at, + writes: AtomicUsize::new(0), + release: tokio::sync::Notify::new(), + }; + let started = AtomicUsize::new(0); + let finished = AtomicUsize::new(0); + let second_started = tokio::sync::Notify::new(); + let mut manifest = pending_manifest(4); + let result = tokio::time::timeout( + std::time::Duration::from_secs(5), + export_data_concurrent( + &storage, + &mut manifest, + 2, + &crate::data::progress::NoopProgress, + |id, _| { + let (started, finished, second_started, storage) = + (&started, &finished, &second_started, &storage); + async move { + started.fetch_add(1, Ordering::SeqCst); + match id { + 1 => second_started.notified().await, + 2 => { + second_started.notify_one(); + storage.release.notified().await; + } + _ => panic!("new chunk started after manifest persistence failed"), + } + finished.fetch_add(1, Ordering::SeqCst); + Ok(vec![format!("data/public/{id}/a.parquet")]) + } + }, + ), + ) + .await + .unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("injected manifest failure") + ); + let expected = if fail_at == 1 { 0 } else { 2 }; + assert_eq!(started.load(Ordering::SeqCst), expected); + assert_eq!(finished.load(Ordering::SeqCst), expected); + } + } + #[test] fn test_next_eligible_chunk_scans_in_order() { let manifest = pending_manifest(3); diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index dce924a717d..9faf0009829 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -30,6 +30,7 @@ pub(super) struct CopyOptions { pub(super) format: DataFormat, pub(super) time_range: TimeRange, pub(super) parallelism: usize, + pub(crate) experimental_metric_export: bool, } pub(super) struct CopyTarget { @@ -210,6 +211,9 @@ pub(crate) async fn execute_copy_database_from( fn build_with_options(options: &CopyOptions) -> String { let mut parts = vec![format!("FORMAT='{}'", options.format)]; + if options.experimental_metric_export { + parts.push("experimental_metric_export='true'".to_string()); + } if let Some(start) = options.time_range.start { parts.push(format!( "START_TIME='{}'", diff --git a/src/cli/src/data/export_v2/error.rs b/src/cli/src/data/export_v2/error.rs index 078f5bd2e1f..dbb4a88ae46 100644 --- a/src/cli/src/data/export_v2/error.rs +++ b/src/cli/src/data/export_v2/error.rs @@ -24,6 +24,14 @@ use snafu::{Location, Snafu}; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Experimental Metric export requires Parquet and a frontend that explicitly reports experimental_metric_export=true. Use an upgraded, consistently configured frontend entry point." + ))] + MetricExportUnavailable { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid URI '{}': {}", uri, reason))] InvalidUri { uri: String, @@ -211,7 +219,8 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { - Error::InvalidUri { .. } + Error::MetricExportUnavailable { .. } + | Error::InvalidUri { .. } | Error::UnsupportedScheme { .. } | Error::SchemaOnlyModeMismatch { .. } | Error::ResumeConfigMismatch { .. } diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index 17da8aea0b0..a5bc02ca7fb 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -40,6 +40,7 @@ use crate::data::export_v2::manifest::{MANIFEST_FILE, Manifest}; #[cfg(test)] use crate::data::export_v2::schema::SchemaDefinition; use crate::data::export_v2::schema::{SCHEMA_DIR, SCHEMAS_FILE, SchemaSnapshot}; +use crate::data::path::data_dir_for_schema_chunk; struct RemoteLocation { bucket_or_container: String, @@ -285,6 +286,21 @@ pub trait SnapshotStorage: Send + Sync { /// Lists files recursively under a relative prefix. async fn list_files_recursive(&self, prefix: &str) -> Result>; + /// Checks fresh output directories or removes a terminated, unfinished chunk's files. + async fn prepare_export_chunk( + &self, + schemas: &[String], + chunk_id: u32, + resume: bool, + ) -> Result<()> { + let _ = (schemas, chunk_id, resume); + InvalidUriSnafu { + uri: "snapshot", + reason: "storage does not support preparing export chunks", + } + .fail() + } + /// Deletes the entire snapshot (for --force). async fn delete_snapshot(&self) -> Result<()>; } @@ -696,6 +712,56 @@ impl SnapshotStorage for OpenDalStorage { Ok(files) } + async fn prepare_export_chunk( + &self, + schemas: &[String], + chunk_id: u32, + resume: bool, + ) -> Result<()> { + let mut files = Vec::new(); + for schema in schemas { + let prefix = data_dir_for_schema_chunk(schema, chunk_id); + let mut entries = match self.object_store.lister_with(&prefix).recursive(true).await { + Ok(entries) => entries, + Err(error) if error.kind() == ErrorKind::NotFound => continue, + Err(error) => { + return Err(error).context(StorageOperationSnafu { + operation: format!("list {prefix}"), + }); + } + }; + while let Some(entry) = entries.try_next().await.context(StorageOperationSnafu { + operation: format!("list {prefix}"), + })? { + let path = entry.path(); + if path == prefix && entry.metadata().is_dir() { + continue; + } + let name = path.strip_prefix(&prefix).unwrap_or(""); + if !resume || entry.metadata().is_dir() || !valid_chunk_filename(name) { + return InvalidUriSnafu { + uri: path, + reason: "expected an empty new chunk or direct Parquet files in an owned unfinished chunk", + }.fail(); + } + files.push(path.to_string()); + } + } + // Validate every schema before deleting any files; COPY starts only after this returns. + for path in files { + match self.object_store.delete(&path).await { + Ok(()) => {} + Err(error) if error.kind() == ErrorKind::NotFound => {} + Err(error) => { + return Err(error).context(StorageOperationSnafu { + operation: format!("delete unfinished chunk file {path}"), + }); + } + } + } + Ok(()) + } + async fn delete_snapshot(&self) -> Result<()> { self.object_store .delete_with("/") @@ -707,6 +773,12 @@ impl SnapshotStorage for OpenDalStorage { } } +fn valid_chunk_filename(name: &str) -> bool { + name.strip_suffix(".parquet") + .is_some_and(|stem| !stem.is_empty()) + && !name.contains(['/', '\\', '\0']) +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -971,6 +1043,140 @@ mod tests { assert!(storage.exists().await.unwrap()); } + #[tokio::test] + async fn test_prepare_export_chunk_preserves_other_chunks_and_schema_artifacts() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + for path in [ + "data/public/1/keep.parquet", + "data/public/2/cpu.v1.parquet", + "data/other/2/a.parquet", + "schema/ddl/public.sql", + ] { + storage.write_text(path, "original").await.unwrap(); + } + let schemas = vec!["public".to_string(), "other".to_string()]; + assert!( + storage + .prepare_export_chunk(&schemas, 2, false) + .await + .is_err() + ); + assert!( + storage + .file_exists("data/public/2/cpu.v1.parquet") + .await + .unwrap() + ); + storage + .prepare_export_chunk(&schemas, 2, true) + .await + .unwrap(); + storage + .prepare_export_chunk(&schemas, 2, true) + .await + .unwrap(); + assert!( + !storage + .file_exists("data/public/2/cpu.v1.parquet") + .await + .unwrap() + ); + assert!(!storage.file_exists("data/other/2/a.parquet").await.unwrap()); + assert_eq!( + storage + .read_text("data/public/1/keep.parquet") + .await + .unwrap(), + "original" + ); + assert_eq!( + storage.read_text("schema/ddl/public.sql").await.unwrap(), + "original" + ); + } + + #[tokio::test] + async fn test_prepare_export_chunk_rejects_foreign_or_nested_entries_before_deletion() { + for foreign in ["notes.txt", "nested/a.parquet", "empty/"] { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + storage + .write_text("data/public/2/cpu.parquet", "keep") + .await + .unwrap(); + let path = format!("data/other/2/{foreign}"); + if foreign.ends_with('/') { + storage.create_dir_all(&path).await.unwrap(); + } else { + storage.write_text(&path, "foreign").await.unwrap(); + } + assert!( + storage + .prepare_export_chunk(&["public".into(), "other".into()], 2, true) + .await + .is_err() + ); + assert_eq!( + storage + .read_text("data/public/2/cpu.parquet") + .await + .unwrap(), + "keep" + ); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn test_prepare_export_chunk_reports_delete_failure() { + use std::os::unix::fs::PermissionsExt; + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + storage + .write_text("data/public/2/a.parquet", "keep") + .await + .unwrap(); + let chunk_dir = dir.path().join("data/public/2"); + std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o555)).unwrap(); + let result = storage + .prepare_export_chunk(&["public".into()], 2, true) + .await; + std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o755)).unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("delete unfinished chunk file") + ); + assert_eq!( + storage.read_text("data/public/2/a.parquet").await.unwrap(), + "keep" + ); + } + + #[tokio::test] + async fn test_prepare_export_chunk_encodes_schema_boundary() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + let schema = "../other".to_string(); + let owned = format!("{}a.parquet", data_dir_for_schema_chunk(&schema, 2)); + storage.write_text(&owned, "remove").await.unwrap(); + storage + .write_text("data/other/2/a.parquet", "keep") + .await + .unwrap(); + storage + .prepare_export_chunk(&[schema], 2, true) + .await + .unwrap(); + assert!(!storage.file_exists(&owned).await.unwrap()); + assert_eq!( + storage.read_text("data/other/2/a.parquet").await.unwrap(), + "keep" + ); + } + #[tokio::test] async fn test_delete_snapshot_only_removes_rooted_contents() { let parent = tempdir().unwrap(); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index cba27129dd6..e295d248f79 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -99,6 +99,18 @@ impl DatabaseClient { /// Execute sql query. pub async fn sql(&self, sql: &str, schema: &str) -> Result>>> { + let body = self.sql_response(sql, schema).await?; + Ok(body.output().first().and_then(|output| match output { + GreptimeQueryOutput::Records(records) => Some(records.rows().clone()), + GreptimeQueryOutput::AffectedRows(_) => None, + })) + } + + pub(crate) async fn sql_response( + &self, + sql: &str, + schema: &str, + ) -> Result { let url = format!("http://{}/v1/sql", self.addr); let params = [ ("db", format!("{}-{}", self.catalog, schema)), @@ -138,11 +150,7 @@ impl DatabaseClient { reason: "cannot get response text".to_string(), })?; - let body = serde_json::from_str::(&text).context(SerdeJsonSnafu)?; - Ok(body.output().first().and_then(|output| match output { - GreptimeQueryOutput::Records(records) => Some(records.rows().clone()), - GreptimeQueryOutput::AffectedRows(_) => None, - })) + serde_json::from_str::(&text).context(SerdeJsonSnafu) } } diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 793ab94205e..86ba2cb9bfb 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -50,6 +50,8 @@ pub struct FrontendOptions { /// even if a request sets the `auto_create_table` hint to `true`. When `true` /// (default), the per-request hint still applies. Default: `true`. pub auto_create_table: bool, + /// Enables experimental Parquet database exports using shared Metric scans. + pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -93,6 +95,7 @@ impl Default for FrontendOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, + experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -273,6 +276,9 @@ max_batch_rows = 25 #[test] fn test_toml() { let opts = FrontendOptions::default(); + assert!(!opts.experimental_metric_export); + let enabled: FrontendOptions = toml::from_str("experimental_metric_export = true").unwrap(); + assert!(enabled.experimental_metric_export); let toml_string = toml::to_string(&opts).unwrap(); assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index e16955f2877..0d1f36ea524 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -122,6 +122,7 @@ lazy_static! { #[derive(Clone)] pub struct Instance { frontend_peer_addr: String, + experimental_metric_export: bool, catalog_manager: CatalogManagerRef, pipeline_operator: Arc, statement_executor: Arc, @@ -391,6 +392,21 @@ impl Instance { } _ => { query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?; + if let Statement::ShowVariables(show) = &stmt + && show + .variable + .to_string() + .eq_ignore_ascii_case("experimental_metric_export") + { + return self.show_metric_export_capability(); + } + if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(CopyDatabase::To( + arg, + ))) = &stmt + && export_database::parse_metric_export_requested(&arg.with)? + { + return self.copy_metric_database(arg.clone(), query_ctx).await; + } self.statement_executor .execute_sql(stmt, query_ctx) .await @@ -2734,6 +2750,116 @@ mod tests { ) } + struct ReleasedExportSource { + schema: GtSchemaRef, + channels: std::sync::Mutex, oneshot::Receiver<()>)>>, + } + + impl DataSource for ReleasedExportSource { + fn get_stream( + &self, + _request: ScanRequest, + ) -> std::result::Result { + let (started, release) = self.channels.lock().unwrap().take().unwrap(); + let schema = self.schema.clone(); + let stream = futures::stream::once(async move { + started.send(()).unwrap(); + release.await.unwrap(); + Ok(RecordBatch::new_empty(schema)) + }); + Ok(Box::pin(RecordBatchStreamWrapper::new( + self.schema.clone(), + Box::pin(stream), + ))) + } + } + + #[tokio::test] + async fn test_metric_export_http_timeout_drains_ordinary_writer() { + let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let info = test_table_info(1024, "source").unwrap(); + let source = Arc::new(Table::new( + Arc::new(info.clone()), + FilterPushDownType::Unsupported, + Arc::new(ReleasedExportSource { + schema: info.meta.schema.clone(), + channels: std::sync::Mutex::new(Some((started_tx, release_rx))), + }), + )); + let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source); + let kv = Arc::new(MemoryKvBackend::new()); + let instance = FrontendBuilder::new( + FrontendOptions { + experimental_metric_export: true, + ..Default::default() + }, + kv.clone(), + test_cache_registry(kv).unwrap(), + catalog, + Arc::new(client::client_manager::NodeClients::default()), + Arc::new(NoopProcedureExecutor), + Arc::new(ProcessManager::new("export-timeout".into(), None)), + ) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed(destination.path()) + .unwrap(), + ) + .try_build() + .await + .unwrap(); + let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions { + timeout: Duration::from_secs(2), + ..Default::default() + }) + .with_sql_handler(Arc::new(instance)) + .build(); + let app = server.build(server.make_app()).unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server_task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let sql = format!( + "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", + destination.path().display() + ); + let request = tokio::spawn(async move { + reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{addr}/v1/sql")) + .form(&[("sql", sql)]) + .send() + .await + .unwrap() + }); + tokio::time::timeout(Duration::from_secs(5), started_rx) + .await + .unwrap() + .unwrap(); + let response = request.await.unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); + // A dropped ordinary stream would close this receiver before release. + release_tx.send(()).unwrap(); + let file = destination.path().join("source.parquet"); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if std::fs::read(&file) + .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + server_task.abort(); + } + fn pending_table( table_id: u32, table_name: &str, diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 7a6ba474168..01e9ef034e3 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -387,6 +387,7 @@ impl FrontendBuilder { Ok(Instance { frontend_peer_addr, + experimental_metric_export: self.options.experimental_metric_export, catalog_manager: self.catalog_manager, pipeline_operator, statement_executor, diff --git a/src/frontend/src/instance/export_database.rs b/src/frontend/src/instance/export_database.rs index 6594490548c..f36bcf562c7 100644 --- a/src/frontend/src/instance/export_database.rs +++ b/src/frontend/src/instance/export_database.rs @@ -12,14 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use auth::{ PermissionAction, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, PermissionTableTargets, }; +use common_error::ext::BoxedError; +use common_query::Output; +use common_recordbatch::RecordBatches; +use common_telemetry::error; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::StringVector; use operator::statement::export_database::{DatabaseExportSummary, PreparedDatabaseExport}; use session::context::QueryContextRef; use snafu::ResultExt; use sql::ast::{Ident, ObjectName}; +use sql::statements::OptionMap; use sql::statements::copy::{Copy, CopyDatabase, CopyDatabaseArgument}; use sql::statements::statement::Statement; use table::requests::CopyDatabaseRequest; @@ -28,7 +38,65 @@ use tokio_util::sync::CancellationToken; use crate::error::{PermissionSnafu, Result}; use crate::instance::Instance; +pub(crate) fn parse_metric_export_requested(options: &OptionMap) -> Result { + match options.get("experimental_metric_export") { + None | Some("false") => Ok(false), + Some("true") => Ok(true), + Some(_) => Err(operator::error::InvalidDatabaseExportSnafu { + reason: "experimental_metric_export must be true or false", + } + .build() + .into()), + } +} + impl Instance { + pub(crate) fn show_metric_export_capability(&self) -> Result { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "EXPERIMENTAL_METRIC_EXPORT", + ConcreteDataType::string_datatype(), + false, + )])); + let batches = RecordBatches::try_from_columns( + schema, + vec![Arc::new(StringVector::from(vec![ + self.experimental_metric_export.to_string(), + ])) as datatypes::vectors::VectorRef], + ) + .map_err(BoxedError::new) + .context(crate::error::ExternalSnafu)?; + Ok(Output::new_with_record_batches(batches)) + } + + pub(crate) async fn copy_metric_database( + &self, + arg: CopyDatabaseArgument, + ctx: QueryContextRef, + ) -> Result { + if !self.experimental_metric_export { + return Err(operator::error::InvalidDatabaseExportSnafu { + reason: "experimental_metric_export is disabled on this frontend", + } + .build() + .into()); + } + let req = operator::statement::to_copy_database_request(arg, &ctx)?; + let plan = self.prepare_database_export(req, None, &ctx).await?; + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let executor = self.statement_executor.clone(); + // Dropping the request only cancels; the runtime task retains and drains started I/O. + let task = common_runtime::spawn_query(async move { + let result = executor.export_database(plan, &cancellation, ctx).await; + if let Err(err) = &result { + error!(err; "Experimental database export failed after draining started work"); + } + result + }); + let summary = task.await.context(operator::error::JoinTaskSnafu)??; + Ok(Output::new_with_affected_rows(summary.rows)) + } + #[allow(dead_code)] async fn export_database( &self, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 61804b30d4c..6822dc97b75 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -870,7 +870,7 @@ fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result< /// Converts [CopyDatabaseArgument] to [CopyDatabaseRequest]. /// This function extracts the necessary info including catalog/database name, time range, etc. -fn to_copy_database_request( +pub fn to_copy_database_request( arg: CopyDatabaseArgument, query_ctx: &QueryContextRef, ) -> Result { diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 4884b5ee091..80ba4cf11a5 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -44,6 +44,8 @@ pub struct StandaloneOptions { /// Upper bound: when `false`, missing tables are never auto-created even if a /// request sets the `auto_create_table` hint to `true`. Default: `true`. pub auto_create_table: bool, + /// Enables experimental Parquet database exports using shared Metric scans. + pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -91,6 +93,7 @@ impl Default for StandaloneOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, + experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -154,6 +157,7 @@ impl StandaloneOptions { FrontendOptions { default_timezone: cloned_opts.default_timezone, auto_create_table: cloned_opts.auto_create_table, + experimental_metric_export: cloned_opts.experimental_metric_export, max_in_flight_write_bytes: cloned_opts.max_in_flight_write_bytes, write_bytes_exhausted_policy: cloned_opts.write_bytes_exhausted_policy, http: cloned_opts.http, @@ -306,6 +310,9 @@ flow_notification_queue_capacity = 17 #[test] fn test_query_options_propagated_to_components() { let mut options = StandaloneOptions::default(); + assert!(!options.frontend_options().experimental_metric_export); + options.experimental_metric_export = true; + assert!(options.frontend_options().experimental_metric_export); options.query.parallelism = 4; assert_eq!(options.frontend_options().query.parallelism, 4); diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index d097432b828..b5603bd29df 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,6 +29,7 @@ cache.workspace = true catalog.workspace = true chrono.workspace = true clap.workspace = true +cli.workspace = true client.workspace = true cmd.workspace = true common-base.workspace = true diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index e875f2d7a01..307a7f6daa2 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -87,6 +87,7 @@ pub struct GreptimeDbStandaloneBuilder { slow_query_options: SlowQueryOptions, event_recorder_options: EventRecorderOptions, auto_create_table: bool, + experimental_metric_export: bool, } impl GreptimeDbStandaloneBuilder { @@ -106,9 +107,17 @@ impl GreptimeDbStandaloneBuilder { }, event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, + experimental_metric_export: false, } } + /// Enables experimental Metric export for the standalone test instance. + #[must_use] + pub fn with_experimental_metric_export(mut self) -> Self { + self.experimental_metric_export = true; + self + } + #[must_use] pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self { self.auto_create_table = auto_create_table; @@ -367,6 +376,7 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + experimental_metric_export: self.experimental_metric_export, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { experimental_enable_resource_info: true, diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 8d5ea3e42bd..e87bf387979 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -746,3 +746,487 @@ async fn database_export_preserves_valid_table_names() { } } } + +#[derive(clap::Parser)] +struct ExportDataCli { + #[command(subcommand)] + command: cli::DataCommand, +} + +async fn run_data_cli(args: &[&str]) -> Result<(), common_error::ext::BoxedError> { + use clap::Parser; + ExportDataCli::try_parse_from(std::iter::once("greptime-data").chain(args.iter().copied())) + .unwrap() + .command + .build() + .await? + .do_work() + .await +} + +async fn export_http(instance: Arc) -> (String, tokio::task::JoinHandle<()>) { + let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions::default()) + .with_sql_handler(instance) + .build(); + let app = server.build(server.make_app()).unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (addr, task) +} + +struct FailSecondChunk { + fail: std::sync::atomic::AtomicBool, + copies: std::sync::atomic::AtomicUsize, +} + +impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { + type Error = frontend::error::Error; + + fn pre_execute( + &self, + statement: Option<&sql::statements::statement::Statement>, + _plan: Option<&datafusion_expr::LogicalPlan>, + _ctx: session::context::QueryContextRef, + ) -> Result<(), Self::Error> { + use sql::statements::copy::{Copy, CopyDatabase}; + use sql::statements::statement::Statement; + if let Some(Statement::Copy(Copy::CopyDatabase(CopyDatabase::To(arg)))) = statement { + self.copies + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if arg.location.ends_with("/z_later/2/") + && self.fail.swap(false, std::sync::atomic::Ordering::SeqCst) + { + return Err(operator::error::InvalidDatabaseExportSnafu { + reason: "injected terminated COPY failure", + } + .build() + .into()); + } + } + Ok(()) + } +} + +#[tokio::test] +async fn metric_export_v2_cli_resume_roundtrip() { + metric_export_v2_cli_roundtrip(false).await; +} + +#[tokio::test] +#[ignore = "requires GT_S3_* credentials and an S3-compatible test bucket"] +async fn metric_export_v2_cli_s3_resume_roundtrip() { + metric_export_v2_cli_roundtrip(true).await; +} + +async fn metric_export_v2_cli_roundtrip(s3: bool) { + use servers::interceptor::SqlQueryInterceptorRef; + let plugins = common_base::Plugins::new(); + let faults = Arc::new(FailSecondChunk { + fail: std::sync::atomic::AtomicBool::new(true), + copies: std::sync::atomic::AtomicUsize::new(0), + }); + plugins.insert::>(faults.clone()); + let source = GreptimeDbStandaloneBuilder::new("metric_v2_source") + .with_experimental_metric_export() + .with_plugin(plugins) + .build() + .await; + let instance = source.fe_instance(); + let mut names = Vec::new(); + for physical in ["v2_a", "v2_b"] { + let (logical, _, renamed) = + create_metric_export_source_tables(instance, physical, "dense").await; + sql( + instance, + &format!("ALTER TABLE {renamed} RENAME {physical}"), + ) + .await; + names.extend(logical); + names.push(format!("{physical}_excluded")); + } + sql( + instance, + "CREATE TABLE audit (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX)", + ) + .await; + sql(instance, "INSERT INTO audit VALUES ('a',1,1),('z',NULL,3)").await; + sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; + sql(instance, "CREATE DATABASE z_later").await; + sql( + instance, + "CREATE TABLE z_later.events (ts TIMESTAMP TIME INDEX, val BIGINT)", + ) + .await; + sql(instance, "INSERT INTO z_later.events VALUES (1,10),(3,20)").await; + names.push("audit".into()); + if s3 { + sql( + instance, + "CREATE TABLE bulk (host STRING PRIMARY KEY, payload STRING, ts TIMESTAMP TIME INDEX)", + ) + .await; + for batch in 0..64 { + let rows = (0..64) + .map(|row| { + let payload = (0..128) + .map(|_| uuid::Uuid::new_v4().simple().to_string()) + .collect::(); + format!("('{}','{payload}',1)", batch * 64 + row) + }) + .collect::>() + .join(","); + sql(instance, &format!("INSERT INTO bulk VALUES {rows}")).await; + } + names.push("bulk".into()); + } + let (addr, server) = export_http(instance.clone()).await; + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let (uri, store, storage_args) = if s3 { + let endpoint = std::env::var("GT_S3_ENDPOINT_URL").unwrap(); + let bucket = std::env::var("GT_S3_BUCKET").unwrap(); + let region = std::env::var("GT_S3_REGION").unwrap(); + let key = std::env::var("GT_S3_ACCESS_KEY_ID").unwrap(); + let secret = std::env::var("GT_S3_ACCESS_KEY").unwrap(); + let root = format!("pr04b-{}", uuid::Uuid::new_v4()); + let store = object_store::ObjectStore::new( + object_store::services::S3::default() + .endpoint(&endpoint) + .bucket(&bucket) + .root(&root) + .region(®ion) + .access_key_id(&key) + .secret_access_key(&secret), + ) + .unwrap(); + ( + format!("s3://{bucket}/{root}"), + store, + vec![ + "--s3".into(), + "--s3-endpoint".into(), + endpoint, + "--s3-region".into(), + region, + "--s3-access-key-id".into(), + key, + "--s3-secret-access-key".into(), + secret, + ], + ) + } else { + ( + url::Url::from_file_path(destination.path()) + .unwrap() + .to_string(), + object_store::ObjectStore::new( + object_store::services::Fs::default().root(destination.path().to_str().unwrap()), + ) + .unwrap(), + Vec::::new(), + ) + }; + let mut args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--schemas", + "public,z_later", + "--experimental-metric-export", + "--no-proxy", + "--start-time", + "1970-01-01T00:00:00Z", + "--end-time", + "1970-01-01T00:00:00.006Z", + "--chunk-time-window", + "3ms", + "--progress", + "never", + ]; + args.extend(storage_args.iter().map(String::as_str)); + assert!(run_data_cli(&args).await.is_err()); + let before: cli::export_v2::manifest::Manifest = + serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); + assert_eq!( + before.chunks[0].status, + cli::export_v2::manifest::ChunkStatus::Completed + ); + assert_eq!( + before.chunks[1].status, + cli::export_v2::manifest::ChunkStatus::Failed + ); + if s3 { + assert!( + store + .stat("data/public/1/bulk.parquet") + .await + .unwrap() + .content_length() + > 5 * 1024 * 1024 + ); + } + let mut preserved = Vec::new(); + for path in &before.chunks[0].files { + preserved.push((path.clone(), store.read(path).await.unwrap().to_vec())); + } + assert!(store.exists("data/public/2/audit.parquet").await.unwrap()); + store + .write("data/public/2/unknown.txt", "keep") + .await + .unwrap(); + let copies = faults.copies.load(std::sync::atomic::Ordering::SeqCst); + assert!(run_data_cli(&args).await.is_err()); + assert_eq!( + faults.copies.load(std::sync::atomic::Ordering::SeqCst), + copies + ); + assert_eq!( + store + .read("data/public/2/unknown.txt") + .await + .unwrap() + .to_vec(), + b"keep" + ); + store.delete("data/public/2/unknown.txt").await.unwrap(); + run_data_cli(&args).await.unwrap(); + let after: cli::export_v2::manifest::Manifest = + serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); + assert!(after.is_complete()); + assert_eq!( + serde_json::to_value(&before.chunks[0]).unwrap(), + serde_json::to_value(&after.chunks[0]).unwrap() + ); + for (path, bytes) in preserved { + assert_eq!(bytes, store.read(&path).await.unwrap().to_vec()); + } + let mut verify = vec!["export-v2", "verify", "--snapshot", &uri]; + verify.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&verify).await.unwrap(); + let target = GreptimeDbStandaloneBuilder::new("metric_v2_target") + .build() + .await; + for id in 0..24 { + sql( + target.fe_instance(), + &format!("CREATE TABLE occupied_{id} (ts TIMESTAMP TIME INDEX)"), + ) + .await; + } + let (target_addr, target_server) = export_http(target.fe_instance().clone()).await; + let state = destination.path().join("restore-state.json"); + let mut import = vec![ + "import-v2", + "--addr", + &target_addr, + "--from", + &uri, + "--no-proxy", + "--state-path", + state.to_str().unwrap(), + "--progress", + "never", + ]; + import.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&import).await.unwrap(); + for name in names { + let source_table = table(instance, &name).await; + let target_table = table(target.fe_instance(), &name).await; + assert_ne!( + source_table.table_info().table_id(), + target_table.table_info().table_id() + ); + assert_eq!( + source_table.schema().column_schemas(), + target_table.schema().column_schemas() + ); + let query = format!("SELECT * FROM \"{name}\" ORDER BY ts,host"); + assert_eq!( + values(instance, &query).await, + values(target.fe_instance(), &query).await + ); + } + assert_eq!( + values(instance, "SELECT * FROM z_later.events ORDER BY ts").await, + values( + target.fe_instance(), + "SELECT * FROM z_later.events ORDER BY ts" + ) + .await + ); + server.abort(); + target_server.abort(); + if s3 { + store.delete_with("/").recursive(true).await.unwrap(); + } +} + +#[tokio::test] +async fn metric_export_v2_disabled_and_legacy_cli() { + let source = GreptimeDbStandaloneBuilder::new("metric_v2_disabled") + .build() + .await; + let (addr, server) = export_http(source.fe_instance().clone()).await; + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let uri = url::Url::from_file_path(destination.path()) + .unwrap() + .to_string(); + let manifest = destination.path().join("manifest.json"); + std::fs::write(&manifest, b"preserve-before-capability-check").unwrap(); + let args = [ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--force", + "--experimental-metric-export", + "--no-proxy", + "--progress", + "never", + ]; + assert!(run_data_cli(&args).await.is_err()); + assert_eq!( + std::fs::read(&manifest).unwrap(), + b"preserve-before-capability-check" + ); + std::fs::remove_file(&manifest).unwrap(); + sql( + source.fe_instance(), + "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, val BIGINT)", + ) + .await; + sql(source.fe_instance(), "INSERT INTO ordinary VALUES (1,2)").await; + run_data_cli(&[ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--schemas", + "public", + "--no-proxy", + "--progress", + "never", + ]) + .await + .unwrap(); + run_data_cli(&["export-v2", "verify", "--snapshot", &uri]) + .await + .unwrap(); + server.abort(); +} + +#[tokio::test] +async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() { + use axum::http::StatusCode; + use serde_json::json; + let column = json!({"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "String"}); + let records = |rows, columns| { + json!({"records": { + "schema": {"column_schemas": columns}, + "rows": rows + }}) + }; + let valid = records(json!([["true"]]), json!([column.clone()])); + let mut responses = [ + (StatusCode::BAD_REQUEST, json!([])), + (StatusCode::OK, json!([])), + (StatusCode::OK, json!([[true]])), + (StatusCode::OK, json!([["true", "extra"]])), + (StatusCode::OK, json!([["true"], ["true"]])), + ] + .map(|(status, rows)| (status, json!([records(rows, json!([column.clone()]))]))) + .to_vec(); + responses.extend([ + ( + StatusCode::OK, + json!([ + valid.clone(), + records(json!([["false"]]), json!([column.clone()])) + ]), + ), + ( + StatusCode::OK, + json!([records(json!([["true"]]), json!([]))]), + ), + ( + StatusCode::OK, + json!([records(json!([["true"]]), json!([column.clone(), column]))]), + ), + ( + StatusCode::OK, + json!([records( + json!([["true"]]), + json!([{"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "Boolean"}]) + )]), + ), + ( + StatusCode::OK, + json!([records( + json!([["true"]]), + json!([{"name": "OTHER", "data_type": "String"}]) + )]), + ), + ]); + for (status, output) in responses { + let app = + axum::Router::new().route( + "/v1/sql", + axum::routing::post( + move |axum::Form(form): axum::Form< + std::collections::HashMap, + >| async move { + assert_eq!(form["sql"], "SHOW VARIABLES experimental_metric_export"); + ( + status, + axum::Json(json!({ + "execution_time_ms": 0, + "output": output + })), + ) + }, + ), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let destination = tempfile::tempdir().unwrap(); + let snapshot = destination.path().join("snapshot"); + let uri = url::Url::from_file_path(&snapshot).unwrap().to_string(); + let args = [ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--experimental-metric-export", + "--force", + "--no-proxy", + "--progress", + "never", + ]; + let error = run_data_cli(&args).await.unwrap_err(); + if status == StatusCode::OK { + assert!( + error.to_string().contains("Metric export requires"), + "{error}" + ); + } + assert!(!snapshot.exists()); + std::fs::create_dir(&snapshot).unwrap(); + let manifest = snapshot.join("manifest.json"); + std::fs::write(&manifest, b"preserve").unwrap(); + assert!(run_data_cli(&args).await.is_err()); + assert_eq!(std::fs::read(manifest).unwrap(), b"preserve"); + server.abort(); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1ad39ae1726..e4e30baeba8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2360,6 +2360,7 @@ mem_threshold_on_create = "auto" r#" enable_telemetry = true auto_create_table = true +experimental_metric_export = false max_in_flight_write_bytes = "0KiB" write_bytes_exhausted_policy = "wait" init_regions_in_background = false From edbcb8224ebc96f9eedeb7d43ac60d71f3386cba Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:19:07 +0000 Subject: [PATCH 42/63] fix: fail startup on duplicate region engine configs (#9281) * fix: reject duplicate region engine configurations Signed-off-by: Lei, HUANG * refactor: reuse canonical engine names in config validation Replace duplicated engine-name literals with the existing common-catalog constants so duplicate-config validation uses the shared engine names. Keep the TOML regression inputs independent to verify the public configuration tags. Signed-off-by: Lei, HUANG * fix: clarify zero-based duplicate engine config indices State explicitly that duplicate region engine configuration indices are zero-based so users can map them to the order of TOML entries. Preserve the existing index values and error classification. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- config/config.md | 4 +- config/datanode.example.toml | 1 + config/standalone.example.toml | 1 + src/datanode/src/datanode.rs | 109 +++++++++++++++++++++++++++++++-- src/datanode/src/error.rs | 12 ++++ 5 files changed, 119 insertions(+), 8 deletions(-) diff --git a/config/config.md b/config/config.md index 7385c8efb0d..16384b406f4 100644 --- a/config/config.md +++ b/config/config.md @@ -170,7 +170,7 @@ | `storage.http_client.timeout` | String | `30s` | The total request timeout, applied from when the request starts connecting until the response body has finished.
Also considered a total deadline. | | `storage.http_client.pool_idle_timeout` | String | `90s` | The timeout for idle sockets being kept-alive. | | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | -| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines. | +| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | | `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | @@ -631,7 +631,7 @@ | `storage.http_client.timeout` | String | `30s` | The total request timeout, applied from when the request starts connecting until the response body has finished.
Also considered a total deadline. | | `storage.http_client.pool_idle_timeout` | String | `90s` | The timeout for idle sockets being kept-alive. | | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | -| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines. | +| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | | `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index f99e1e124f9..45baaefbdfb 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -464,6 +464,7 @@ skip_ssl_validation = false # endpoint = "https://storage.googleapis.com" ## The region engine options. You can configure multiple region engines. +## Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. [[region_engine]] ## The Mito engine options. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index cc4be3f5930..e59e7ba1b4c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -674,6 +674,7 @@ skip_ssl_validation = false # endpoint = "https://storage.googleapis.com" ## The region engine options. You can configure multiple region engines. +## Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. [[region_engine]] ## The Mito engine options. diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 15e83fc1c7b..ae4b0723e1a 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -14,11 +14,13 @@ //! Datanode implementation. +use std::collections::HashMap; use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; use common_base::Plugins; +use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE}; use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; @@ -63,9 +65,10 @@ use tokio::sync::Notify; use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, - DataFusionSnafu, GetMetadataSnafu, InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, - MissingNodeIdSnafu, ObjectStoreWalNotSupportedSnafu, OpenLogStoreSnafu, Result, - ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, + DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu, + InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu, + ObjectStoreWalNotSupportedSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, + ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver, @@ -253,6 +256,7 @@ impl DatanodeBuilder { } pub async fn build(mut self) -> Result { + validate_region_engine_config(&self.opts.region_engine)?; let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?; set_default_prefix(self.opts.default_column_prefix.as_deref()) .map_err(BoxedError::new) @@ -710,6 +714,27 @@ impl DatanodeBuilder { } } +/// Rejects repeated engine types before their configs can silently overwrite each other. +fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> { + let mut engine_indices = HashMap::new(); + for (index, config) in configs.iter().enumerate() { + let engine = match config { + RegionEngineConfig::Mito(_) => MITO_ENGINE, + RegionEngineConfig::File(_) => FILE_ENGINE, + RegionEngineConfig::Metric(_) => METRIC_ENGINE, + }; + if let Some(first_index) = engine_indices.insert(engine, index) { + return DuplicateRegionEngineConfigSnafu { + engine, + first_index, + duplicate_index: index, + } + .fail(); + } + } + Ok(()) +} + /// Rejects an object store WAL config the log store cannot run on. fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> { let prefix = config.prefix.trim(); @@ -879,6 +904,7 @@ async fn open_all_regions( mod tests { use std::assert_matches; use std::collections::{BTreeMap, HashMap}; + use std::io::Write; use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -886,6 +912,7 @@ mod tests { use cache::build_datanode_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; + use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_meta::cache::LayeredCacheRegistryBuilder; @@ -893,15 +920,17 @@ mod tests { use common_meta::key::datanode_table::DatanodeTableManager; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; - use common_test_util::temp_dir::create_temp_dir; + use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use mito2::engine::MITO_ENGINE_NAME; use store_api::region_request::RegionRequest; use store_api::storage::RegionId; - use crate::config::DatanodeOptions; - use crate::datanode::{DatanodeBuilder, validate_object_store_wal_config}; + use crate::config::{DatanodeOptions, RegionEngineConfig}; + use crate::datanode::{ + DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, + }; use crate::error::Error; use crate::tests::{MockRegionEngine, mock_region_server}; @@ -1014,6 +1043,74 @@ mod tests { std::fs::read_dir(dir).unwrap().next().is_none() } + #[test] + fn test_validate_unique_region_engine_config() { + let mut opts = DatanodeOptions::default(); + validate_region_engine_config(&opts.region_engine).unwrap(); + + opts.region_engine + .push(RegionEngineConfig::Metric(Default::default())); + validate_region_engine_config(&opts.region_engine).unwrap(); + + // Omitted engines still use their defaults during engine construction. + validate_region_engine_config(&[]).unwrap(); + } + + #[tokio::test] + async fn test_build_rejects_duplicate_region_engine_config() { + for (engine, first_index) in [("mito", 0), ("file", 1), ("metric", 2)] { + let mut file = create_named_temp_file(); + write!( + file, + r#" + [[region_engine]] + [region_engine.mito] + global_write_buffer_size = "4GiB" + min_compaction_interval = "600s" + + [[region_engine]] + [region_engine.file] + + [[region_engine]] + [region_engine.metric] + + [[region_engine]] + [region_engine.{engine}] + "# + ) + .unwrap(); + let opts = DatanodeOptions::load_layered_options( + Some(file.path().to_str().unwrap()), + "DATANODE_DUPLICATE_REGION_ENGINE_UT", + ) + .unwrap(); + let data_home = create_temp_dir("duplicate-region-engine-config"); + let mut builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::default(), + Arc::new(MemoryKvBackend::new()), + ); + builder.opts.region_engine = opts.region_engine; + let err = build_err(builder).await; + + let Error::DuplicateRegionEngineConfig { + engine: actual_engine, + first_index: actual_first_index, + duplicate_index, + .. + } = &err + else { + panic!("unexpected error for {engine}: {err:?}"); + }; + assert_eq!(*actual_engine, engine); + assert_eq!(*actual_first_index, first_index); + assert_eq!(*duplicate_index, 3); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + // Rejected before the builder creates any storage or log store. + assert!(is_empty_dir(data_home.path())); + } + } + #[tokio::test] async fn test_build_rejects_invalid_object_store_wal_config() { let cases = [ diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index c391efffcf0..b62708aa343 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -299,6 +299,17 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Duplicate region engine config '{engine}' at region_engine[{first_index}] and region_engine[{duplicate_index}] (indices are zero-based); each engine type may be configured only once" + ))] + DuplicateRegionEngineConfig { + engine: &'static str, + first_index: usize, + duplicate_index: usize, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid object store WAL config, {} = {:?}: {}", field, value, reason))] InvalidObjectStoreWalConfig { field: &'static str, @@ -486,6 +497,7 @@ impl ErrorExt for Error { | GcConfigMismatch { .. } | ParseAddr { .. } | TomlFormat { .. } + | DuplicateRegionEngineConfig { .. } | InvalidObjectStoreWalConfig { .. } | BuildDatanode { .. } => StatusCode::InvalidArguments, From 19c127d9c3a11262f9652832aaf4d7e41b36bd0e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 00:07:04 +0000 Subject: [PATCH 43/63] feat: restore packed metric snapshots (#9250) * test: cover snapshot parquet restore compatibility Signed-off-by: jeremyhi * feat: restore packed metric snapshots Signed-off-by: jeremyhi * fix: stream large packed parquet entries Signed-off-by: jeremyhi * fix: allow default S3 endpoint in packed restore test Signed-off-by: jeremyhi * fix: skip unconfigured S3 in packed restore test Signed-off-by: jeremyhi * test: use portable file URLs in packed restore fixtures Signed-off-by: jeremyhi * fix: validate packed snapshot structure before restore Signed-off-by: jeremyhi * fix: preserve strict manifest decoding and verify fixture Signed-off-by: jeremyhi * fix: reject packed layout on both database export paths Signed-off-by: jeremyhi * fix: implement file size in coordinator test storage Signed-off-by: jeremyhi * test: use expect_err for rejected export layouts Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 2 + src/cli/Cargo.toml | 2 + src/cli/src/data/export_v2/command.rs | 24 +- src/cli/src/data/export_v2/coordinator.rs | 3 + src/cli/src/data/export_v2/data.rs | 8 +- src/cli/src/data/export_v2/manifest.rs | 170 ++++- src/cli/src/data/import_v2.rs | 1 + src/cli/src/data/import_v2/command.rs | 89 ++- src/cli/src/data/import_v2/error.rs | 13 +- src/cli/src/data/import_v2/packed.rs | 225 ++++++ src/cli/src/data/snapshot_storage.rs | 15 +- src/cli/src/database.rs | 89 +++ src/common/datasource/src/error.rs | 4 + .../datasource/src/file_format/parquet.rs | 2 + .../src/file_format/parquet/packed_reader.rs | 351 ++++++++++ src/common/datasource/src/lib.rs | 1 + src/common/datasource/src/object_store.rs | 13 +- src/common/datasource/src/packed_snapshot.rs | 188 +++++ src/frontend/src/instance.rs | 10 + src/frontend/src/instance/import_packed.rs | 114 +++ src/operator/src/error.rs | 7 +- src/operator/src/statement.rs | 1 + src/operator/src/statement/copy_database.rs | 10 +- src/operator/src/statement/copy_table_from.rs | 99 +++ src/operator/src/statement/database_copy.rs | 12 + src/operator/src/statement/export_database.rs | 2 + src/operator/src/statement/import_packed.rs | 436 ++++++++++++ src/servers/src/http.rs | 38 + tests-integration/Cargo.toml | 1 + .../tests/export_logical_tables.rs | 661 ++++++++++++++++++ .../cases/snapshot_parquet_copy/case.toml | 8 + .../cases/snapshot_parquet_copy/setup.sql | 4 + .../cases/snapshot_parquet_copy/verify.result | 12 + .../cases/snapshot_parquet_copy/verify.sql | 2 + 34 files changed, 2583 insertions(+), 34 deletions(-) create mode 100644 src/cli/src/data/import_v2/packed.rs create mode 100644 src/common/datasource/src/file_format/parquet/packed_reader.rs create mode 100644 src/common/datasource/src/packed_snapshot.rs create mode 100644 src/frontend/src/instance/import_packed.rs create mode 100644 src/operator/src/statement/import_packed.rs create mode 100644 tests/compatibility/cases/snapshot_parquet_copy/case.toml create mode 100644 tests/compatibility/cases/snapshot_parquet_copy/setup.sql create mode 100644 tests/compatibility/cases/snapshot_parquet_copy/verify.result create mode 100644 tests/compatibility/cases/snapshot_parquet_copy/verify.sql diff --git a/Cargo.lock b/Cargo.lock index dc0dcf2561c..7dadf276b7a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2016,6 +2016,7 @@ dependencies = [ "common-base", "common-catalog", "common-config", + "common-datasource", "common-error", "common-grpc", "common-macro", @@ -2049,6 +2050,7 @@ dependencies = [ "servers", "session", "snafu 0.8.6", + "sql", "standalone", "store-api", "table", diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 27aaeed73b5..3a7dc1d68e4 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -28,6 +28,7 @@ client = { workspace = true, features = ["testing"] } common-base.workspace = true common-catalog.workspace = true common-config.workspace = true +common-datasource.workspace = true common-error.workspace = true common-grpc.workspace = true common-macro.workspace = true @@ -62,6 +63,7 @@ serde_json.workspace = true servers.workspace = true session.workspace = true snafu.workspace = true +sql.workspace = true standalone.workspace = true store-api.workspace = true table.workspace = true diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index fc958e9e3e0..e0f67eb796c 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -504,7 +504,7 @@ impl ExportCreate { let mut manifest = self.storage.read_manifest().await?; // Check version compatibility - if manifest.version != MANIFEST_VERSION { + if manifest.version != MANIFEST_VERSION || manifest.data_layout.is_some() { return ManifestVersionMismatchSnafu { expected: MANIFEST_VERSION, found: manifest.version, @@ -1008,7 +1008,9 @@ fn directory_word(count: usize) -> &'static str { } fn snapshot_status(manifest: &Manifest) -> &'static str { - if manifest.schema_only { + if manifest.validate_layout().is_err() { + "unsupported" + } else if manifest.schema_only { "schema-only" } else if manifest.is_complete() { "complete" @@ -1131,11 +1133,17 @@ async fn verify_snapshot(storage: &OpenDalStorage) -> Result { problems: Vec::new(), }; - if report.manifest.version != MANIFEST_VERSION { - report.push_error(format!( - "Manifest version mismatch: expected {}, found {}", - MANIFEST_VERSION, report.manifest.version - )); + if let Err(reason) = report.manifest.validate_layout() { + report.push_error(reason); + } else if report.manifest.is_packed() + && let Err(error) = crate::data::import_v2::packed::validate_snapshot( + storage, + &report.manifest, + &report.manifest.schemas, + ) + .await + { + report.push_error(error.to_string()); } if !report.schema_index_exists { @@ -1432,7 +1440,7 @@ fn safe_manifest_data_file_path(path: &str) -> Option<&str> { fn print_verify_report(snapshot: &str, report: &VerifyReport) { println!("Verifying snapshot: {}", report.manifest.snapshot_id); println!(" Location: {}", snapshot); - if report.manifest.version == MANIFEST_VERSION { + if report.manifest.validate_layout().is_ok() { println!(" Manifest: OK (version {})", report.manifest.version); } else { println!( diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index b70897337a7..5f95c43603e 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -383,6 +383,9 @@ mod tests { async fn read_manifest(&self) -> Result { unreachable!() } + async fn file_size(&self, _: &str) -> Result> { + unreachable!() + } async fn write_schema( &self, _: &crate::data::export_v2::schema::SchemaSnapshot, diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index 9faf0009829..8723154e4b1 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -191,13 +191,19 @@ pub(crate) async fn execute_copy_database_from( schema: &str, source: &CopySource, format: DataFormat, + packed: bool, ) -> Result<()> { let sql = format!( - r#"COPY DATABASE "{}"."{}" FROM '{}' WITH (FORMAT='{}'){};"#, + r#"COPY DATABASE "{}"."{}" FROM '{}' WITH (FORMAT='{}'{}){};"#, escape_sql_identifier(catalog), escape_sql_identifier(schema), escape_sql_literal(&source.location), format, + if packed { + ", metric_data_layout='packed'" + } else { + "" + }, source.connection ); let safe_sql = source.mask_sql(&sql); diff --git a/src/cli/src/data/export_v2/manifest.rs b/src/cli/src/data/export_v2/manifest.rs index 918288bb51d..5dd9d6f1684 100644 --- a/src/cli/src/data/export_v2/manifest.rs +++ b/src/cli/src/data/export_v2/manifest.rs @@ -14,6 +14,7 @@ //! Manifest data structures for Export/Import V2. +use std::collections::HashSet; use std::time::Duration; use std::{fmt, str}; @@ -27,7 +28,7 @@ use crate::data::export_v2::error::{ TimeParseInvalidFormatSnafu, }; -/// Current manifest format version. +/// Manifest format version produced by the current exporter. pub const MANIFEST_VERSION: u32 = 1; /// Manifest file name within snapshot directory. @@ -234,6 +235,8 @@ impl str::FromStr for DataFormat { pub struct Manifest { /// Manifest format version for compatibility checking. pub version: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub data_layout: Option, /// Unique snapshot identifier. pub snapshot_id: Uuid, /// Catalog name. @@ -259,6 +262,81 @@ pub struct Manifest { } impl Manifest { + /// Decodes version-specific required fields without changing version-1 defaults. + pub fn from_json(data: &[u8]) -> serde_json::Result { + use serde::de::Error; + + let value: serde_json::Value = serde_json::from_slice(data)?; + if value["version"] == 2 { + let chunks = value["chunks"].as_array().ok_or_else(|| { + serde_json::Error::custom("version-2 manifest requires a chunks array") + })?; + for chunk in chunks { + if !chunk["files"].is_array() { + return Err(serde_json::Error::custom( + "version-2 chunk requires a files array", + )); + } + } + } + serde_json::from_slice(data) + } + + /// Validates the encoding contract before any snapshot operation. + pub fn validate_layout(&self) -> std::result::Result<(), String> { + match (self.version, self.data_layout.as_deref(), self.format) { + (1, None, _) => Ok(()), + (2, Some(common_datasource::packed_snapshot::PACKED_LAYOUT), DataFormat::Parquet) => { + if self.checksum.is_some() || self.chunks.iter().any(|c| c.checksum.is_some()) { + return Err("version-2 checksums are unsupported".into()); + } + if self.schema_only != self.chunks.is_empty() { + return Err("version-2 schema_only requires empty chunks; data snapshots require chunks".into()); + } + let schemas: HashSet<_> = self.schemas.iter().collect(); + if schemas.len() != self.schemas.len() { + return Err("duplicate version-2 schemas".into()); + } + let mut ids = HashSet::new(); + for chunk in &self.chunks { + if chunk.id == 0 || !ids.insert(chunk.id) { + return Err("version-2 chunk IDs must be unique and positive".into()); + } + let files: HashSet<_> = chunk.files.iter().collect(); + if files.len() != chunk.files.len() { + return Err("duplicate version-2 chunk files".into()); + } + if chunk.status == ChunkStatus::Skipped && !chunk.files.is_empty() { + return Err("skipped version-2 chunks must have no files".into()); + } + for file in &chunk.files { + if file.contains('\\') + || file + .split('/') + .any(|part| part.is_empty() || part == "." || part == "..") + || !self.schemas.iter().any(|schema| { + file.starts_with(&crate::data::path::data_dir_for_schema_chunk( + schema, chunk.id, + )) + }) + { + return Err(format!("invalid version-2 chunk file: {file}")); + } + } + } + Ok(()) + } + _ => Err(format!( + "Manifest version mismatch or unsupported layout: version {}, layout {:?}, format {}", + self.version, self.data_layout, self.format + )), + } + } + + pub fn is_packed(&self) -> bool { + self.version == 2 + } + pub fn new_for_export( catalog: String, schemas: Vec, @@ -293,6 +371,7 @@ impl Manifest { let now = Utc::now(); Self { version: MANIFEST_VERSION, + data_layout: None, snapshot_id: Uuid::new_v4(), catalog, schemas, @@ -316,6 +395,7 @@ impl Manifest { let now = Utc::now(); Self { version: MANIFEST_VERSION, + data_layout: None, snapshot_id: Uuid::new_v4(), catalog, schemas, @@ -413,6 +493,94 @@ fn generate_single_chunk(time_range: &TimeRange) -> Vec { #[cfg(test)] mod tests { + #[test] + fn packed_manifest_read_contract_preserves_legacy_writer() { + let mut manifest = + super::Manifest::new_schema_only("greptime".into(), vec!["public".into()]); + assert!(manifest.validate_layout().is_ok()); + manifest.version = 2; + assert!(manifest.validate_layout().is_err()); + manifest.data_layout = Some("metric-parquet-packs".into()); + assert!(manifest.validate_layout().is_ok()); + manifest.format = super::DataFormat::Csv; + assert!(manifest.validate_layout().is_err()); + manifest.format = super::DataFormat::Parquet; + manifest.version = 3; + assert!(manifest.validate_layout().is_err()); + manifest.version = 1; + assert!(manifest.validate_layout().is_err()); + } + + #[test] + fn packed_required_fields_and_checksums_preserve_v1_defaults() { + let mut manifest = + super::Manifest::new_schema_only("greptime".into(), vec!["public".into()]); + manifest + .chunks + .push(super::ChunkMeta::new(1, super::TimeRange::unbounded())); + let legacy = serde_json::to_value(manifest).unwrap(); + for version in [1, 2] { + let mut value = legacy.clone(); + value["version"] = version.into(); + if version == 2 { + value["data_layout"] = "metric-parquet-packs".into(); + value["schema_only"] = false.into(); + } + let raw = serde_json::to_string(&value).unwrap(); + for duplicate in [ + format!( + "{{\"checksum\":\"unsupported\",\"checksum\":null,{}", + &raw[1..] + ), + format!("{{\"chunks\":[],{}", &raw[1..]), + raw.replacen("\"files\":[]", "\"files\":[],\"files\":[]", 1), + raw.replacen( + "\"files\":[]", + "\"checksum\":\"unsupported\",\"checksum\":null,\"files\":[]", + 1, + ), + ] { + assert!(super::Manifest::from_json(duplicate.as_bytes()).is_err()); + } + for field in ["chunks", "files"] { + let mut missing = value.clone(); + if field == "chunks" { + missing.as_object_mut().unwrap().remove(field); + } else { + missing["chunks"][0].as_object_mut().unwrap().remove(field); + } + assert_eq!( + super::Manifest::from_json(&serde_json::to_vec(&missing).unwrap()).is_ok(), + version == 1 + ); + } + for checksum in [ + None, + Some(serde_json::Value::Null), + Some("".into()), + Some("sha256:abc".into()), + ] { + for chunk_level in [false, true] { + let mut candidate = value.clone(); + if let Some(checksum) = &checksum { + if chunk_level { + candidate["chunks"][0]["checksum"] = checksum.clone(); + } else { + candidate["checksum"] = checksum.clone(); + } + } + let decoded = + super::Manifest::from_json(&serde_json::to_vec(&candidate).unwrap()) + .unwrap(); + assert_eq!( + decoded.validate_layout().is_ok(), + version == 1 || checksum.as_ref().is_none_or(serde_json::Value::is_null) + ); + } + } + } + } + use std::time::Duration; use chrono::{TimeZone, Utc}; diff --git a/src/cli/src/data/import_v2.rs b/src/cli/src/data/import_v2.rs index b197f59b3b4..c508b943e35 100644 --- a/src/cli/src/data/import_v2.rs +++ b/src/cli/src/data/import_v2.rs @@ -38,6 +38,7 @@ mod command; pub(crate) mod coordinator; pub mod error; pub mod executor; +pub(crate) mod packed; pub(crate) mod state; pub use command::ImportV2Command; diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index d885c4fc781..33cf9cb1ffb 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -27,15 +27,15 @@ use snafu::{OptionExt, ResultExt}; use crate::Tool; use crate::common::ObjectStoreConfig; use crate::data::export_v2::data::{build_copy_source, execute_copy_database_from}; -use crate::data::export_v2::manifest::{ChunkMeta, ChunkStatus, DataFormat, MANIFEST_VERSION}; +use crate::data::export_v2::manifest::{ChunkMeta, ChunkStatus, DataFormat}; use crate::data::import_v2::coordinator::{ ImportResumeConfig, ImportTaskExecutor, build_import_tasks, chunk_has_schema_files, import_with_resume_session_with_progress, prepare_import_resume, }; use crate::data::import_v2::error::{ ChunkImportFailedSnafu, EmptyChunkManifestSnafu, ImportStatePathUnavailableSnafu, - IncompleteSnapshotSnafu, ManifestVersionMismatchSnafu, MissingChunkDataSnafu, Result, - SchemaNotInSnapshotSnafu, SnapshotStorageSnafu, + IncompleteSnapshotSnafu, MissingChunkDataSnafu, Result, SchemaNotInSnapshotSnafu, + SnapshotStorageSnafu, }; use crate::data::import_v2::executor::{DdlExecutor, DdlStatement}; use crate::data::import_v2::state::{ImportTaskKey, default_state_path}; @@ -219,13 +219,14 @@ impl Import { manifest.snapshot_id, manifest.version, manifest.schema_only ); - // Check version compatibility - if manifest.version != MANIFEST_VERSION { - return ManifestVersionMismatchSnafu { - expected: MANIFEST_VERSION, - found: manifest.version, - } - .fail(); + manifest.validate_layout().map_err(|reason| { + crate::data::import_v2::error::InvalidPackedSnapshotSnafu { reason }.build() + })?; + if manifest.is_packed() { + self.database_client + .require_packed_import() + .await + .context(crate::data::import_v2::error::DatabaseSnafu)?; } info!("Snapshot contains {} schema(s)", manifest.schemas.len()); @@ -251,6 +252,15 @@ impl Import { Vec::new() }; + if manifest.is_packed() { + crate::data::import_v2::packed::validate_snapshot( + self.storage.as_ref(), + &manifest, + &schemas_to_import, + ) + .await?; + } + // 4. Dry-run mode: print DDL and exit if self.dry_run { info!("Dry-run mode - DDL statements to execute:"); @@ -317,6 +327,7 @@ impl Import { let executor = CopyDatabaseImportTaskExecutor { import: self, format: manifest.format, + packed: manifest.is_packed(), }; let progress = build_progress_reporter(self.progress); import_with_resume_session_with_progress(resume_session, &executor, progress.as_ref()) @@ -358,6 +369,7 @@ impl Import { struct CopyDatabaseImportTaskExecutor<'a> { import: &'a Import, format: DataFormat, + packed: bool, } #[async_trait] @@ -380,6 +392,7 @@ impl ImportTaskExecutor for CopyDatabaseImportTaskExecutor<'_> { &task.schema, &source, self.format, + self.packed, ) .await .context(ChunkImportFailedSnafu { @@ -684,6 +697,13 @@ mod tests { Ok(self.manifest.clone()) } + async fn file_size( + &self, + _path: &str, + ) -> crate::data::export_v2::error::Result> { + unimplemented!("not needed in import_v2::command tests") + } + async fn write_manifest( &self, _manifest: &Manifest, @@ -731,6 +751,55 @@ mod tests { } } + #[tokio::test] + async fn packed_capability_failure_precedes_ddl_reads_and_state_writes() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let n = socket.read(&mut request).await.unwrap(); + assert!(String::from_utf8_lossy(&request[..n]).starts_with("GET /v1/capabilities ")); + socket + .write_all( + b"HTTP/1.1 404 Not Found\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + ) + .await + .unwrap(); + }); + let dir = tempfile::tempdir().unwrap(); + let state_path = dir.path().join("state.json"); + let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]); + manifest.version = 2; + manifest.data_layout = Some("metric-parquet-packs".into()); + let import = Import { + catalog: "greptime".into(), + schemas: None, + dry_run: false, + progress: ProgressMode::default(), + task_parallelism: 1, + state_path: Some(state_path.clone()), + snapshot_uri: "file:///unused".into(), + storage_config: Default::default(), + storage: Box::new(StubStorage { + manifest, + files_by_prefix: Default::default(), + }), + database_client: DatabaseClient::new( + address.to_string(), + "greptime".into(), + None, + std::time::Duration::from_secs(5), + None, + true, + ), + }; + assert!(import.run().await.is_err()); + assert!(!state_path.exists()); + server.await.unwrap(); + } + fn parse_command(extra: &[&str]) -> ImportV2Command { let mut args = vec![ "import-v2", diff --git a/src/cli/src/data/import_v2/error.rs b/src/cli/src/data/import_v2/error.rs index 967a62bd831..d1699634448 100644 --- a/src/cli/src/data/import_v2/error.rs +++ b/src/cli/src/data/import_v2/error.rs @@ -25,6 +25,9 @@ use crate::data::export_v2::manifest::ChunkStatus; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Invalid packed snapshot: {reason}"))] + InvalidPackedSnapshot { reason: String }, + #[snafu(display("Snapshot not found at '{}'", uri))] SnapshotNotFound { uri: String, @@ -32,14 +35,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Manifest version mismatch: expected {}, found {}", expected, found))] - ManifestVersionMismatch { - expected: u32, - found: u32, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Schema '{}' not found in snapshot", schema))] SchemaNotInSnapshot { schema: String, @@ -190,9 +185,9 @@ impl ErrorExt for Error { match self { Error::SnapshotNotFound { .. } | Error::SchemaNotInSnapshot { .. } - | Error::ManifestVersionMismatch { .. } | Error::IncompleteSnapshot { .. } | Error::EmptyChunkManifest { .. } + | Error::InvalidPackedSnapshot { .. } | Error::MissingChunkData { .. } => StatusCode::InvalidArguments, Error::ImportStatePathUnavailable { .. } | Error::ImportStateUnknownTask { .. } diff --git a/src/cli/src/data/import_v2/packed.rs b/src/cli/src/data/import_v2/packed.rs new file mode 100644 index 00000000000..303c866ba69 --- /dev/null +++ b/src/cli/src/data/import_v2/packed.rs @@ -0,0 +1,225 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; + +use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; +use snafu::ResultExt; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::statement::Statement; + +use crate::data::export_v2::manifest::{ChunkStatus, Manifest}; +use crate::data::import_v2::error::{InvalidPackedSnapshotSnafu, Result, SnapshotStorageSnafu}; +use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; +use crate::data::snapshot_storage::SnapshotStorage; + +/// Matches every completed schema chunk to its source DDL before target writes. +pub(crate) async fn validate_snapshot( + storage: &dyn SnapshotStorage, + manifest: &Manifest, + schemas: &[String], +) -> Result<()> { + manifest + .validate_layout() + .map_err(|reason| InvalidPackedSnapshotSnafu { reason }.build())?; + if manifest.schema_only + && !storage + .list_files_recursive("data/") + .await + .context(SnapshotStorageSnafu)? + .is_empty() + { + return InvalidPackedSnapshotSnafu { + reason: "schema-only snapshot contains data objects", + } + .fail(); + } + for schema in schemas { + let ddl = storage + .read_text(&ddl_path_for_schema(schema)) + .await + .context(SnapshotStorageSnafu)?; + let statements = ParserContext::create_with_dialect( + &ddl, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .map_err(|e| { + InvalidPackedSnapshotSnafu { + reason: e.to_string(), + } + .build() + })?; + let mut names = HashSet::new(); + for statement in statements { + if let Statement::CreateTable(create) = statement { + if create.engine.eq_ignore_ascii_case("metric") + && create.options.get("on_physical_table").is_none() + { + continue; + } + let parts: Vec<_> = create + .name + .0 + .iter() + .map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let allowed = match parts.as_slice() { + [Some(_)] => true, + [Some(s), Some(_)] => *s == schema, + [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, + _ => false, + }; + if !allowed { + return InvalidPackedSnapshotSnafu { + reason: "snapshot table DDL escapes its catalog/schema", + } + .fail(); + } + let name = parts.last().and_then(|p| *p).unwrap_or_default(); + if !names.insert(name.to_string()) { + return InvalidPackedSnapshotSnafu { + reason: "duplicate snapshot table DDL", + } + .fail(); + } + } + } + for chunk in &manifest.chunks { + if chunk.status == ChunkStatus::Skipped && !names.is_empty() { + return InvalidPackedSnapshotSnafu { + reason: "packed chunks require zero-row streams for empty tables", + } + .fail(); + } + if chunk.status != ChunkStatus::Completed { + continue; + } + let path = format!( + "{}{PACK_INDEX_FILE}", + data_dir_for_schema_chunk(schema, chunk.id) + ); + let text = storage + .read_text(&path) + .await + .context(SnapshotStorageSnafu)?; + let index: PackIndex = serde_json::from_str(&text).map_err(|e| { + InvalidPackedSnapshotSnafu { + reason: e.to_string(), + } + .build() + })?; + let prefix = data_dir_for_schema_chunk(schema, chunk.id); + let expected: HashSet<_> = std::iter::once(path) + .chain(index.objects.iter().map(|o| format!("{prefix}{}", o.path))) + .collect(); + let actual: HashSet<_> = chunk + .files + .iter() + .filter(|file| file.starts_with(&prefix)) + .cloned() + .collect(); + if actual != expected { + return InvalidPackedSnapshotSnafu { + reason: "chunk manifest inventory differs from indexed objects", + } + .fail(); + } + index + .validate_membership(names.iter().map(String::as_str)) + .map_err(|e| { + InvalidPackedSnapshotSnafu { + reason: e.to_string(), + } + .build() + })?; + for object in &index.objects { + let path = format!("{prefix}{}", object.path); + let length = storage + .file_size(&path) + .await + .context(SnapshotStorageSnafu)?; + if length != Some(object.length) { + return InvalidPackedSnapshotSnafu { + reason: format!( + "object {path}: expected length {}, actual {length:?}", + object.length + ), + } + .fail(); + } + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data::export_v2::manifest::{ChunkMeta, TimeRange}; + use crate::data::snapshot_storage::OpenDalStorage; + + #[tokio::test] + async fn packed_membership_uses_literal_names_and_excludes_physical_and_views() { + let dir = tempfile::Builder::new() + .prefix("packed snapshot # ") + .tempdir() + .unwrap(); + let uri = url::Url::from_file_path(dir.path()).unwrap(); + let storage = OpenDalStorage::from_uri(uri.as_str(), &Default::default()).unwrap(); + storage.write_text("schema/ddl/public.sql", "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\";").await.unwrap(); + let path = "data/public/1/pack-index.json"; + let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":12}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0}]}); + storage.write_text(path, &index.to_string()).await.unwrap(); + assert!(dir.path().join(path).is_file()); + storage + .write_text("data/public/1/pack-0.bin", "abcdefghijkl") + .await + .unwrap(); + let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]); + let mut chunk = ChunkMeta::new(1, TimeRange::unbounded()); + chunk.mark_completed(vec![path.into(), "data/public/1/pack-0.bin".into()], None); + manifest.version = 2; + manifest.data_layout = Some("metric-parquet-packs".into()); + manifest.schema_only = false; + manifest.chunks.push(chunk); + validate_snapshot(&storage, &manifest, &manifest.schemas) + .await + .unwrap(); + let mut missing = index.clone(); + missing["tables"] = serde_json::json!([]); + storage + .write_text(path, &missing.to_string()) + .await + .unwrap(); + assert!( + validate_snapshot(&storage, &manifest, &manifest.schemas) + .await + .is_err() + ); + let mut escaped = index; + escaped["objects"][0]["path"] = "../pack-0.bin".into(); + storage + .write_text(path, &escaped.to_string()) + .await + .unwrap(); + assert!( + validate_snapshot(&storage, &manifest, &manifest.schemas) + .await + .is_err() + ); + } +} diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index a5bc02ca7fb..b3a5fca84bc 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -268,6 +268,9 @@ pub trait SnapshotStorage: Send + Sync { /// Reads the manifest file. async fn read_manifest(&self) -> Result; + /// Returns file length, or None for a missing path or directory. + async fn file_size(&self, path: &str) -> Result>; + /// Writes the manifest file. async fn write_manifest(&self, manifest: &Manifest) -> Result<()>; @@ -669,7 +672,17 @@ impl SnapshotStorage for OpenDalStorage { ensure_snapshot_exists(self).await?; let data = self.read_file(MANIFEST_FILE).await?; - serde_json::from_slice(&data).context(ManifestParseSnafu) + Manifest::from_json(&data).context(ManifestParseSnafu) + } + + async fn file_size(&self, path: &str) -> Result> { + match self.object_store.stat(path).await { + Ok(metadata) => Ok(metadata.is_file().then(|| metadata.content_length())), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(None), + Err(error) => Err(error).with_context(|_| StorageOperationSnafu { + operation: format!("stat {path}"), + }), + } } async fn write_manifest(&self, manifest: &Manifest) -> Result<()> { diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index e295d248f79..b771ab48d05 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -97,6 +97,54 @@ impl DatabaseClient { self.sql(sql, DEFAULT_SCHEMA_NAME).await } + /// Requires the explicit packed-import protocol before any restore mutation. + pub async fn require_packed_import(&self) -> Result<()> { + let url = format!("http://{}/v1/capabilities", self.addr); + let mut builder = reqwest::Client::builder().timeout(self.timeout); + if let Some(proxy) = self.proxy.clone() { + builder = builder.proxy(proxy); + } + if self.no_proxy { + builder = builder.no_proxy(); + } + let client = builder.build().context(BuildClientSnafu)?; + let mut request = client.get(&url); + if let Some(auth) = &self.auth_header { + request = request.header("Authorization", auth); + } + let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { + reason: "packed import capability request failed", + })?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return crate::error::InvalidArgumentsSnafu { + msg: "target does not support packed import", + } + .fail(); + } + let response = response + .error_for_status() + .with_context(|_| HttpQuerySqlSnafu { + reason: "packed import capability request rejected", + })?; + let body = response.text().await.with_context(|_| HttpQuerySqlSnafu { + reason: "cannot read capability response", + })?; + let value: Value = serde_json::from_str(&body).context(SerdeJsonSnafu)?; + if !value.is_object() { + return crate::error::UnexpectedSnafu { + msg: "invalid capability response: expected JSON object", + } + .fail(); + } + if value.get("metric_packed_import").and_then(Value::as_u64) != Some(1) { + return crate::error::InvalidArgumentsSnafu { + msg: "target does not support packed import version 1", + } + .fail(); + } + Ok(()) + } + /// Execute sql query. pub async fn sql(&self, sql: &str, schema: &str) -> Result>>> { let body = self.sql_response(sql, schema).await?; @@ -170,6 +218,47 @@ pub(crate) fn split_database(database: &str) -> Result<(String, Option)> #[cfg(test)] mod tests { + #[tokio::test] + async fn packed_capability_probe_checks_auth_and_protocol_version() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + for (status, body, supported) in [ + (200, r#"{"metric_packed_import":1,"future":2}"#, true), + (200, "{}", false), + (200, r#"{"metric_packed_import":2}"#, false), + (404, "{}", false), + (401, "{}", false), + (403, "{}", false), + (200, "[]", false), + (200, "invalid-json", false), + ] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let n = socket.read(&mut request).await.unwrap(); + let request = String::from_utf8_lossy(&request[..n]).to_lowercase(); + assert!(request.starts_with("get /v1/capabilities ")); + assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza==")); + socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + }); + let client = super::DatabaseClient::new( + address.to_string(), + "greptime".into(), + Some("user:password".into()), + std::time::Duration::from_secs(5), + None, + true, + ); + assert_eq!( + client.require_packed_import().await.is_ok(), + supported, + "{status}: {body}" + ); + server.await.unwrap(); + } + } + use super::*; #[test] diff --git a/src/common/datasource/src/error.rs b/src/common/datasource/src/error.rs index 1cc04d5a3de..f98b43202f2 100644 --- a/src/common/datasource/src/error.rs +++ b/src/common/datasource/src/error.rs @@ -27,6 +27,9 @@ use url::ParseError; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Invalid packed snapshot: {reason}"))] + InvalidPackedSnapshot { reason: String }, + #[snafu(display("Parquet write cancelled"))] ParquetWriteCancelled {}, @@ -285,6 +288,7 @@ impl ErrorExt for Error { use Error::*; match self { ParquetWriteCancelled {} => StatusCode::Cancelled, + InvalidPackedSnapshot { .. } => StatusCode::InvalidArguments, InvalidParquetWriterLimits {} => StatusCode::InvalidArguments, ParquetWriterResource { .. } => StatusCode::Suspended, BuildBackend { .. } diff --git a/src/common/datasource/src/file_format/parquet.rs b/src/common/datasource/src/file_format/parquet.rs index 0032de6b31e..b8f9420ba73 100644 --- a/src/common/datasource/src/file_format/parquet.rs +++ b/src/common/datasource/src/file_format/parquet.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod packed_reader; + use std::result; use std::sync::Arc; diff --git a/src/common/datasource/src/file_format/parquet/packed_reader.rs b/src/common/datasource/src/file_format/parquet/packed_reader.rs new file mode 100644 index 00000000000..9136b1ff940 --- /dev/null +++ b/src/common/datasource/src/file_format/parquet/packed_reader.rs @@ -0,0 +1,351 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Request-owned, bounded backing buffers for independent Parquet streams. + +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::future::BoxFuture; +use object_store::ObjectStore; +use parquet::arrow::arrow_reader::ArrowReaderOptions; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +use tokio::sync::Mutex; + +pub const WINDOW_SIZE: usize = 8 * 1024 * 1024; +const WINDOW_BUDGET: usize = 2 * WINDOW_SIZE; + +struct Window { + path: String, + offset: u64, + bytes: Bytes, +} + +/// Shared only by table decoders belonging to one authorized COPY request. +/// The lock also coalesces concurrent fetches for the same bytes. +pub struct PackReadWindows { + store: ObjectStore, + windows: Mutex>, +} + +impl PackReadWindows { + pub fn new(store: ObjectStore) -> Arc { + Arc::new(Self { + store, + windows: Mutex::new(Vec::new()), + }) + } + + async fn read(&self, path: &str, object_length: u64, range: Range) -> Result { + if range.start > range.end || range.end > object_length { + return Err(ParquetError::General("packed read outside object".into())); + } + if range.is_empty() { + return Ok(Bytes::new()); + } + let mut windows = self.windows.lock().await; + for window in windows.iter() { + if window.path == path + && range.start >= window.offset + && range.end <= window.offset + window.bytes.len() as u64 + { + return Ok(window.bytes.slice( + (range.start - window.offset) as usize..(range.end - window.offset) as usize, + )); + } + } + let required = usize::try_from(range.end - range.start) + .map_err(|_| ParquetError::General("packed read length overflow".into()))?; + if required > WINDOW_SIZE { + return Err(ParquetError::General( + "packed cache range exceeds window".into(), + )); + } + let fetch_len = (object_length - range.start).min(WINDOW_SIZE as u64) as usize; + // Evict only unpinned allocations. Bytes slices held by a decoder still + // count against the request budget, including their entire backing window. + let mut retained: usize = windows.iter().map(|w| w.bytes.len()).sum(); + while retained.saturating_add(fetch_len) > WINDOW_BUDGET || windows.len() >= 2 { + let Some(i) = windows.iter().position(|w| w.bytes.is_unique()) else { + return Err(ParquetError::General( + "packed read backing-buffer budget exhausted".into(), + )); + }; + retained -= windows.remove(i).bytes.len(); + } + let bytes = self + .store + .read_with(path) + .range(range.start..range.start + fetch_len as u64) + .await + .map_err(|e| ParquetError::External(Box::new(e)))? + .to_bytes(); + if bytes.len() != fetch_len { + return Err(ParquetError::General("short packed object read".into())); + } + // Backend buffers can be slices of larger allocations. Retain an exact + // window allocation; transport staging is separate from cached backing. + let bytes = Bytes::copy_from_slice(&bytes); + common_telemetry::debug!( + path, + offset = range.start, + length = fetch_len, + "Fetched packed read window" + ); + let result = bytes.slice(..required); + windows.push(Window { + path: path.into(), + offset: range.start, + bytes, + }); + Ok(result) + } +} + +/// Pins one complete small stream; larger streams use the streaming reader. +pub struct PackedParquetReader { + windows: Arc, + path: String, + object_length: u64, + offset: u64, + length: u64, + bytes: Option, +} + +impl PackedParquetReader { + pub fn new( + windows: Arc, + path: String, + object_length: u64, + offset: u64, + length: u64, + ) -> Result { + if length < 12 + || length > WINDOW_SIZE as u64 + || offset + .checked_add(length) + .is_none_or(|end| end > object_length) + { + return Err(ParquetError::General("invalid packed Parquet range".into())); + } + Ok(Self { + windows, + path, + object_length, + offset, + length, + bytes: None, + }) + } +} + +impl AsyncFileReader for PackedParquetReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + Box::pin(async move { + if range.start > range.end || range.end > self.length { + return Err(ParquetError::General( + "read outside indexed Parquet stream".into(), + )); + } + if self.bytes.is_none() { + self.bytes = Some( + self.windows + .read( + &self.path, + self.object_length, + self.offset..self.offset + self.length, + ) + .await?, + ); + } + Ok(self + .bytes + .as_ref() + .unwrap() + .slice(range.start as usize..range.end as usize)) + }) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let length = self.length; + let metadata = ParquetMetaDataReader::new() + .with_metadata_options(options.map(|o| o.metadata_options().clone())) + .load_and_finish(self, length) + .await?; + Ok(Arc::new(metadata)) + }) + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{ArrayRef, Int64Array, StringArray}; + use arrow::record_batch::RecordBatch; + use arrow_schema::{Field, Schema}; + use futures::TryStreamExt; + use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}; + + use super::*; + + fn parquet(array: ArrayRef) -> (Vec, RecordBatch) { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + array.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + (writer.into_inner().unwrap(), batch) + } + + #[tokio::test] + async fn independent_schemas_empty_streams_and_shared_ranges() { + let dir = common_test_util::temp_dir::create_temp_dir("packed-reader"); + let store = crate::test_util::test_store(dir.path().to_str().unwrap()); + let fixtures = [ + parquet(Arc::new(Int64Array::from(vec![42]))), + parquet(Arc::new(StringArray::from(vec!["different type"]))), + parquet(Arc::new(Int64Array::from(Vec::::new()))), + ]; + let packed: Vec = fixtures + .iter() + .flat_map(|(b, _)| b.iter().copied()) + .collect(); + store.write("pack-0.bin", packed.clone()).await.unwrap(); + let windows = PackReadWindows::new(store.clone()); + // A pinned shared window serves all the independent stream readers. + let pinned = windows + .read("pack-0.bin", packed.len() as u64, 0..packed.len() as u64) + .await + .unwrap(); + store.delete("pack-0.bin").await.unwrap(); + let mut offset = 0; + for (bytes, expected) in fixtures { + let reader = PackedParquetReader::new( + windows.clone(), + "pack-0.bin".into(), + packed.len() as u64, + offset, + bytes.len() as u64, + ) + .unwrap(); + let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); + assert_eq!(builder.schema(), &expected.schema()); + let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap(); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + expected.num_rows() + ); + if expected.num_rows() > 0 { + assert_eq!(batches[0], expected); + } + offset += bytes.len() as u64; + } + assert_eq!(windows.windows.lock().await.len(), 1); + assert_eq!(pinned.len(), packed.len()); + let isolated = PackReadWindows::new(store); + assert!( + isolated + .read("pack-0.bin", packed.len() as u64, 0..1) + .await + .is_err() + ); + } + + #[tokio::test] + async fn small_stream_pins_one_window_for_multiple_columns() { + use parquet::file::properties::WriterProperties; + let schema = Arc::new(Schema::new( + (0..3) + .map(|i| Field::new(format!("v{i}"), arrow_schema::DataType::Utf8, false)) + .collect::>(), + )); + let value = "x".repeat(1024 * 1024); + let arrays = (0..3) + .map(|_| Arc::new(StringArray::from(vec![value.as_str()])) as ArrayRef) + .collect(); + let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap(); + let mut writer = ArrowWriter::try_new( + Vec::new(), + schema, + Some( + WriterProperties::builder() + .set_dictionary_enabled(false) + .build(), + ), + ) + .unwrap(); + writer.write(&batch).unwrap(); + let bytes = writer.into_inner().unwrap(); + assert!(bytes.len() < WINDOW_SIZE); + let dir = common_test_util::temp_dir::create_temp_dir("packed-columns"); + let store = crate::test_util::test_store(dir.path().to_str().unwrap()); + let mut object = vec![0; WINDOW_SIZE - 1024]; + let offset = object.len() as u64; + object.extend_from_slice(&bytes); + let length = object.len() as u64; + store.write("pack-0.bin", object).await.unwrap(); + let windows = PackReadWindows::new(store); + let pinned = windows.read("pack-0.bin", length, 0..1).await.unwrap(); + let mut reader = PackedParquetReader::new( + windows.clone(), + "pack-0.bin".into(), + length, + offset, + bytes.len() as u64, + ) + .unwrap(); + let metadata = reader.get_metadata(None).await.unwrap(); + let mut columns = Vec::new(); + for column in metadata.row_group(0).columns() { + let (start, len) = column.byte_range(); + columns.push(reader.get_bytes(start..start + len).await.unwrap()); + } + let batches: Vec<_> = ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches, vec![batch]); + let retained = windows.windows.lock().await; + assert_eq!(retained.len(), 2); + assert!(retained.iter().all(|w| w.bytes.len() <= WINDOW_SIZE)); + assert!(retained.iter().map(|w| w.bytes.len()).sum::() <= WINDOW_BUDGET); + assert_eq!(columns.len(), 3); + assert_eq!(pinned[0], 0); + } + + #[tokio::test] + async fn corrupt_metadata_short_reads_and_relative_bounds() { + let dir = common_test_util::temp_dir::create_temp_dir("packed-corrupt"); + let store = crate::test_util::test_store(dir.path().to_str().unwrap()); + store.write("pack-0.bin", vec![0u8; 12]).await.unwrap(); + let windows = PackReadWindows::new(store); + assert!(windows.read("pack-0.bin", 20, 0..20).await.is_err()); + let mut reader = PackedParquetReader::new(windows, "pack-0.bin".into(), 12, 0, 12).unwrap(); + assert!(reader.get_bytes(0..13).await.is_err()); + assert!(reader.get_metadata(None).await.is_err()); + } +} diff --git a/src/common/datasource/src/lib.rs b/src/common/datasource/src/lib.rs index f4c7fdf1202..c832831da80 100644 --- a/src/common/datasource/src/lib.rs +++ b/src/common/datasource/src/lib.rs @@ -19,6 +19,7 @@ pub mod error; pub mod file_format; pub mod lister; pub mod object_store; +pub mod packed_snapshot; pub mod parquet_writer; pub mod share_buffer; #[cfg(test)] diff --git a/src/common/datasource/src/object_store.rs b/src/common/datasource/src/object_store.rs index 37dd11aafdb..ef24f878e50 100644 --- a/src/common/datasource/src/object_store.rs +++ b/src/common/datasource/src/object_store.rs @@ -56,8 +56,17 @@ pub struct BuiltBackend { impl BuiltBackend { /// Checks an explicit input using the same file-type filter as directory listing. pub async fn is_file(&self, path: &str) -> object_store::Result { - // Stat first preserves errors for dangling or inaccessible symlink targets. - if self.object_store.stat(path).await?.mode() != object_store::EntryMode::FILE { + let meta = self.object_store.stat(path).await?; + self.is_file_with_mode(path, meta.mode()).await + } + + /// Applies local file-type protections to a previously fetched object stat. + pub async fn is_file_with_mode( + &self, + path: &str, + mode: object_store::EntryMode, + ) -> object_store::Result { + if mode != object_store::EntryMode::FILE { return Ok(false); } if let Some(root) = &self.local_root { diff --git a/src/common/datasource/src/packed_snapshot.rs b/src/common/datasource/src/packed_snapshot.rs new file mode 100644 index 00000000000..e045b2229a8 --- /dev/null +++ b/src/common/datasource/src/packed_snapshot.rs @@ -0,0 +1,188 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shared on-disk index for one packed snapshot schema chunk. + +use std::collections::{HashMap, HashSet}; + +use serde::{Deserialize, Serialize}; + +use crate::error::{InvalidPackedSnapshotSnafu, Result}; + +pub const PACK_INDEX_FILE: &str = "pack-index.json"; +pub const PACKED_LAYOUT: &str = "metric-parquet-packs"; + +/// Versioned index. Object paths are generated direct children of the chunk. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackIndex { + pub version: u32, + pub objects: Vec, + pub tables: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackObject { + pub path: String, + pub kind: ObjectKind, + pub length: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ObjectKind { + Pack, + Parquet, +} + +/// One complete independent Parquet stream, including its footer. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackTable { + pub table_name: String, + pub object: String, + pub offset: u64, + pub length: u64, + pub row_count: u64, +} + +impl PackIndex { + /// Validates all references and complete, nonoverlapping object coverage. + pub fn validate(&self) -> Result<()> { + let invalid = |reason: &str| InvalidPackedSnapshotSnafu { reason }.build(); + if self.version != 1 { + return Err(invalid("unsupported index version")); + } + let mut objects = HashMap::new(); + for object in &self.objects { + let (prefix, suffix) = match object.kind { + ObjectKind::Pack => ("pack-", ".bin"), + ObjectKind::Parquet => ("table-", ".parquet"), + }; + let generated = object + .path + .strip_prefix(prefix) + .and_then(|s| s.strip_suffix(suffix)) + .is_some_and(|s| !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())); + if !generated + || object.length == 0 + || objects.insert(object.path.as_str(), object).is_some() + { + return Err(invalid("invalid or duplicate object path/length")); + } + } + let mut names = HashSet::new(); + let mut ranges: HashMap<&str, Vec<&PackTable>> = HashMap::new(); + for table in &self.tables { + if table.table_name.is_empty() || !names.insert(table.table_name.as_str()) { + return Err(invalid("empty or duplicate table name")); + } + let object = objects + .get(table.object.as_str()) + .ok_or_else(|| invalid("unknown object reference"))?; + let end = table + .offset + .checked_add(table.length) + .ok_or_else(|| invalid("table range overflow"))?; + if table.length < 12 || end > object.length { + return Err(invalid("invalid Parquet stream range")); + } + ranges.entry(&table.object).or_default().push(table); + } + for object in &self.objects { + let entries = ranges + .get_mut(object.path.as_str()) + .ok_or_else(|| invalid("unreferenced object"))?; + entries.sort_unstable_by_key(|t| t.offset); + if object.kind == ObjectKind::Parquet && entries.len() != 1 { + return Err(invalid("standalone object must contain one table")); + } + let mut end = 0; + for entry in entries { + if entry.offset != end { + return Err(invalid("object ranges have a gap or overlap")); + } + end += entry.length; // Checked against object length above. + } + if end != object.length { + return Err(invalid("object ranges do not cover the full object")); + } + } + Ok(()) + } + + /// Requires exactly the data-bearing tables selected by snapshot DDL. + pub fn validate_membership<'a>(&self, tables: impl IntoIterator) -> Result<()> { + self.validate()?; + let expected: HashSet<_> = tables.into_iter().collect(); + let actual: HashSet<_> = self.tables.iter().map(|t| t.table_name.as_str()).collect(); + if expected != actual { + return InvalidPackedSnapshotSnafu { + reason: "index table membership differs from snapshot DDL", + } + .fail(); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn index() -> PackIndex { + PackIndex { + version: 1, + objects: vec![PackObject { + path: "pack-000000.bin".into(), + kind: ObjectKind::Pack, + length: 24, + }], + tables: ["literal.name", "other"] + .into_iter() + .enumerate() + .map(|(i, name)| PackTable { + table_name: name.into(), + object: "pack-000000.bin".into(), + offset: i as u64 * 12, + length: 12, + row_count: 0, + }) + .collect(), + } + } + + #[test] + fn validates_complete_membership_and_ranges() { + index() + .validate_membership(["other", "literal.name"]) + .unwrap(); + assert!(index().validate_membership(["other"]).is_err()); + for mutate in [ + |i: &mut PackIndex| i.version = 2, + |i: &mut PackIndex| i.objects[0].path = "../pack-0.bin".into(), + |i: &mut PackIndex| i.objects.push(i.objects[0].clone()), + |i: &mut PackIndex| i.tables[1].table_name = "literal.name".into(), + |i: &mut PackIndex| i.tables[1].object = "missing".into(), + |i: &mut PackIndex| i.tables[1].offset = 11, + |i: &mut PackIndex| i.tables[1].length = u64::MAX, + |i: &mut PackIndex| i.tables.pop().map(|_| ()).unwrap(), + ] { + let mut bad = index(); + mutate(&mut bad); + assert!(bad.validate().is_err(), "{bad:?}"); + } + } +} diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 0d1f36ea524..22d2afaa7bc 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -17,6 +17,7 @@ mod dashboard; mod entity_graph; mod export_database; mod grpc; +mod import_packed; mod influxdb; mod jaeger; mod log_handler; @@ -392,6 +393,15 @@ impl Instance { } _ => { query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?; + if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase( + CopyDatabase::From(arg), + )) = &stmt + && arg.with.get("metric_data_layout").is_some() + { + return self + .copy_packed_database(arg.clone(), &stmt, query_ctx) + .await; + } if let Statement::ShowVariables(show) = &stmt && show .variable diff --git a/src/frontend/src/instance/import_packed.rs b/src/frontend/src/instance/import_packed.rs new file mode 100644 index 00000000000..d2cfd825975 --- /dev/null +++ b/src/frontend/src/instance/import_packed.rs @@ -0,0 +1,114 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use auth::{PermissionReq, PermissionTableTarget, PermissionTableTargets}; +use common_query::Output; +use session::context::QueryContextRef; +use snafu::ResultExt; +use sql::statements::copy::CopyDatabaseArgument; +use sql::statements::statement::Statement; +use tokio_util::sync::CancellationToken; + +use crate::error::{PermissionSnafu, Result}; +use crate::instance::Instance; + +impl Instance { + pub(crate) async fn copy_packed_database( + &self, + arg: CopyDatabaseArgument, + stmt: &Statement, + ctx: QueryContextRef, + ) -> Result { + let req = operator::statement::to_copy_database_request(arg, &ctx)?; + let plan = self + .statement_executor + .prepare_packed_import(req, &ctx) + .await?; + let targets = PermissionTableTargets::resolved( + plan.tables + .iter() + .map(|t| { + let info = t.table_info(); + PermissionTableTarget::new(&info.catalog_name, &info.schema_name, &info.name) + }) + .collect(), + ); + self.check_table_permission(&ctx, PermissionReq::SqlStatement(stmt), targets) + .context(PermissionSnafu)?; + let executor = self.statement_executor.clone(); + run_packed_import(move |cancellation| async move { + executor.import_packed(plan, &cancellation, ctx).await + }) + .await + } +} + +async fn run_packed_import(worker: F) -> Result +where + F: FnOnce(CancellationToken) -> Fut, + Fut: std::future::Future> + Send + 'static, +{ + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + // Transport timeout/drop signals cancellation; the worker retains ownership + // until its started inserts finish, even after the response is gone. + common_runtime::spawn_query(worker(cancellation)) + .await + .context(operator::error::JoinTaskSnafu)? + .map_err(Into::into) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::oneshot; + + use super::*; + + #[tokio::test] + async fn packed_request_drop_keeps_worker_alive_until_cleanup() { + let (started_tx, started_rx) = oneshot::channel(); + let (cancelled_tx, cancelled_rx) = oneshot::channel(); + let (drain_tx, drain_rx) = oneshot::channel(); + let (finished_tx, mut finished_rx) = oneshot::channel(); + let mut request = Box::pin(run_packed_import(move |cancellation| async move { + started_tx.send(()).unwrap(); + cancellation.cancelled().await; + cancelled_tx.send(()).unwrap(); + // Model an already-started insert that must finish before worker exit. + drain_rx.await.unwrap(); + finished_tx.send(()).unwrap(); + Ok(Output::new_with_affected_rows(1)) + })); + tokio::select! { + _ = started_rx => {} + _ = &mut request => panic!("worker must wait for cancellation"), + } + drop(request); + tokio::time::timeout(Duration::from_secs(5), cancelled_rx) + .await + .unwrap() + .unwrap(); + assert!(matches!( + finished_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + drain_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(5), finished_rx) + .await + .unwrap() + .unwrap(); + } +} diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 40f968bc413..cd8b8284c64 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -92,6 +92,9 @@ pub enum Error { #[snafu(display("Database export cancelled"))] DatabaseExportCancelled {}, + #[snafu(display("Packed import cancelled"))] + PackedImportCancelled {}, + #[snafu(display("Invalid logical table export: {reason}"))] InvalidLogicalTableExport { reason: String }, @@ -1105,7 +1108,9 @@ impl ErrorExt for Error { StatusCode::InvalidArguments } Error::InvalidDatabaseExport { .. } => StatusCode::InvalidArguments, - Error::DatabaseExportCancelled { .. } => StatusCode::Cancelled, + Error::DatabaseExportCancelled { .. } | Error::PackedImportCancelled { .. } => { + StatusCode::Cancelled + } Error::InvalidLogicalTableExport { .. } => StatusCode::InvalidArguments, Error::LogicalTableExportResource { .. } => StatusCode::Suspended, Error::LogicalTableExportCancelled { .. } => StatusCode::Cancelled, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 6822dc97b75..17d29ab4259 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -25,6 +25,7 @@ mod describe; mod dml; pub mod export_database; pub mod export_logical_tables; +pub mod import_packed; mod kill; pub mod semantic_graph; mod set; diff --git a/src/operator/src/statement/copy_database.rs b/src/operator/src/statement/copy_database.rs index 7a9dd847177..994394cdf3b 100644 --- a/src/operator/src/statement/copy_database.rs +++ b/src/operator/src/statement/copy_database.rs @@ -33,7 +33,7 @@ use crate::error; use crate::statement::StatementExecutor; use crate::statement::database_copy::{ DatabaseExportFile, database_import_source, parse_parallelism_from_option_map, - validate_database_directory, + validate_database_directory, validate_database_export_layout, }; pub(crate) const COPY_DATABASE_TIME_START_KEY: &str = "start_time"; @@ -47,6 +47,7 @@ impl StatementExecutor { req: CopyDatabaseRequest, ctx: QueryContextRef, ) -> error::Result { + validate_database_export_layout(&req.with)?; validate_database_directory(&req.location)?; build_backend_for_write(&req.location, &req.connection, &self.local_file_access) .await @@ -113,6 +114,13 @@ impl StatementExecutor { req: CopyDatabaseRequest, ctx: QueryContextRef, ) -> error::Result { + if let Some(layout) = req.with.get("metric_data_layout") { + return error::InvalidCopyParameterSnafu { + key: "metric_data_layout", + value: layout, + } + .fail(); + } validate_database_directory(&req.location)?; let parallelism = parse_parallelism_from_option_map(&req.with); diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index c69d0477dbd..71efb7266b0 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -374,6 +374,105 @@ impl StatementExecutor { } } + /// Imports one indexed stream with the ordinary COPY schema mapping and inserter. + pub(crate) async fn copy_indexed_parquet< + R: datafusion::parquet::arrow::async_reader::AsyncFileReader + Send + Unpin + 'static, + >( + &self, + reader: R, + table: table::TableRef, + expected_rows: u64, + pending: &mut crate::statement::import_packed::PendingPackedInserts, + cancellation: &tokio_util::sync::CancellationToken, + query_ctx: QueryContextRef, + ) -> Result<()> { + let builder = tokio::select! { + _ = cancellation.cancelled() => return error::PackedImportCancelledSnafu.fail(), + result = ParquetRecordBatchStreamBuilder::new(reader) => result.context(error::ReadParquetMetadataSnafu)?, + }; + let table_schema = table.schema().arrow_schema().clone(); + let (file_projection, table_projection, _) = + generated_schema_projection_and_compatible_file_schema(builder.schema(), &table_schema); + let file_schema = Arc::new( + builder + .schema() + .project(&file_projection) + .context(error::ProjectSchemaSnafu)?, + ); + let target_schema = Arc::new( + table_schema + .project(&table_projection) + .context(error::ProjectSchemaSnafu)?, + ); + ensure_schema_compatible(&file_schema, &target_schema)?; + let stream = builder + .with_batch_size(DEFAULT_BATCH_SIZE) + .build() + .context(error::BuildParquetRecordBatchStreamSnafu)?; + let mut stream = Box::pin(RecordBatchStreamTypeAdapter::new( + target_schema.clone(), + stream, + Some(file_projection), + )); + let info = table.table_info(); + let mut rows = 0u64; + loop { + pending.before_decode().await?; + if cancellation.is_cancelled() { + return error::PackedImportCancelledSnafu.fail(); + } + let batch = tokio::select! { + _ = cancellation.cancelled() => return error::PackedImportCancelledSnafu.fail(), + batch = stream.next() => batch, + }; + let Some(batch) = batch else { + break; + }; + let batch = batch.context(error::ReadDfRecordBatchSnafu)?; + rows += batch.num_rows() as u64; + if rows > expected_rows { + return error::InvalidCopyParameterSnafu { + key: "row_count", + value: expected_rows.to_string(), + } + .fail(); + } + let vectors = Helper::try_into_vectors(batch.columns()).context(IntoVectorsSnafu)?; + let retained_bytes = vectors.iter().map(|v| v.memory_size()).sum::(); + let columns_values = target_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .zip(vectors) + .collect(); + let bytes = retained_bytes.max(batch.get_array_memory_size()); + let inserter = self.inserter.clone(); + let request = InsertRequest { + catalog_name: info.catalog_name.clone(), + schema_name: info.schema_name.clone(), + table_name: info.name.clone(), + columns_values, + skip_wal: query_ctx.skip_wal(), + }; + let ctx = query_ctx.clone(); + pending + .admit( + bytes, + async move { inserter.handle_table_insert(request, ctx).await }, + cancellation, + ) + .await?; + } + if rows != expected_rows { + return error::InvalidCopyParameterSnafu { + key: "row_count", + value: expected_rows.to_string(), + } + .fail(); + } + Ok(()) + } + #[tracing::instrument(skip_all)] pub async fn copy_table_from( &self, diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs index 94ff1df11b5..117104b9d8f 100644 --- a/src/operator/src/statement/database_copy.rs +++ b/src/operator/src/statement/database_copy.rs @@ -59,6 +59,18 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap) -> Result<()> { + if let Some(layout) = options.get("metric_data_layout") { + return error::InvalidCopyParameterSnafu { + key: "metric_data_layout", + value: layout, + } + .fail(); + } + Ok(()) +} + pub(crate) fn validate_database_directory(location: &str) -> Result<()> { ensure!( is_directory_location(location), diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index c9c698a356b..7c5de554487 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -35,6 +35,7 @@ use crate::error::{self, InvalidDatabaseExportSnafu, Result}; use crate::statement::StatementExecutor; use crate::statement::database_copy::{ DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory, + validate_database_export_layout, }; use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; @@ -68,6 +69,7 @@ impl StatementExecutor { req: CopyDatabaseRequest, tables: Vec, ) -> Result { + validate_database_export_layout(&req.with)?; validate_database_directory(&req.location)?; let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; ensure!( diff --git a/src/operator/src/statement/import_packed.rs b/src/operator/src/statement/import_packed.rs new file mode 100644 index 00000000000..7f31804c70b --- /dev/null +++ b/src/operator/src/statement/import_packed.rs @@ -0,0 +1,436 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use common_datasource::file_format::Format; +use common_datasource::file_format::parquet::packed_reader::{ + PackReadWindows, PackedParquetReader, WINDOW_SIZE, +}; +use common_datasource::object_store::{BuiltBackend, build_backend_with_path}; +use common_datasource::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex}; +use common_error::ext::BoxedError; +use common_query::Output; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt}; +use session::context::QueryContextRef; +use snafu::ResultExt; +use table::TableRef; +use table::requests::CopyDatabaseRequest; +use tokio_util::compat::FuturesAsyncReadCompatExt; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::statement::StatementExecutor; +use crate::statement::database_copy::validate_database_directory; + +/// All table identities are resolved before the frontend authorizes any write. +pub struct PreparedPackedImport { + backend: BuiltBackend, + index: PackIndex, + parallelism: usize, + pub tables: Vec, +} + +impl StatementExecutor { + pub async fn prepare_packed_import( + &self, + req: CopyDatabaseRequest, + ctx: &QueryContextRef, + ) -> Result { + validate_database_directory(&req.location)?; + if req.with.get("metric_data_layout").map(String::as_str) != Some("packed") + || !matches!( + Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?, + Format::Parquet(_) + ) + || req + .with + .get("continue_on_error") + .is_some_and(|v| v != "false") + { + return error::InvalidCopyParameterSnafu { + key: "metric_data_layout", + value: "packed requires parquet and continue_on_error=false", + } + .fail(); + } + let backend = + build_backend_with_path(&req.location, &req.connection, &self.local_file_access) + .await + .context(error::BuildBackendSnafu)?; + if !backend + .is_file(PACK_INDEX_FILE) + .await + .context(error::ReadObjectSnafu { + path: PACK_INDEX_FILE, + })? + { + return error::InvalidCopyParameterSnafu { + key: "index", + value: "not a regular file", + } + .fail(); + } + let bytes = backend + .object_store + .read(PACK_INDEX_FILE) + .await + .context(error::ReadObjectSnafu { + path: PACK_INDEX_FILE, + })? + .to_bytes(); + let mut index: PackIndex = serde_json::from_slice(&bytes).map_err(|e| { + error::InvalidCopyParameterSnafu { + key: "pack_index", + value: e.to_string(), + } + .build() + })?; + index + .validate() + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + index + .tables + .sort_by(|a, b| (&a.object, a.offset).cmp(&(&b.object, b.offset))); + let names: Vec<_> = index.tables.iter().map(|t| t.table_name.clone()).collect(); + let mut tables = self + .capture_database_export_tables(&req, Some(&names), ctx) + .await?; + if tables.len() != names.len() { + return error::InvalidCopyParameterSnafu { + key: "index", + value: "physical tables and views are not data targets", + } + .fail(); + } + let positions: HashMap<_, _> = names + .iter() + .enumerate() + .map(|(i, n)| (n.as_str(), i)) + .collect(); + tables.sort_by_key(|table| positions.get(table.table_info().name.as_str()).copied()); + for object in &index.objects { + let meta = backend + .object_store + .stat(&object.path) + .await + .context(error::ReadObjectSnafu { path: &object.path })?; + if !backend + .is_file_with_mode(&object.path, meta.mode()) + .await + .context(error::ReadObjectSnafu { path: &object.path })? + { + return error::InvalidCopyParameterSnafu { + key: "object", + value: &object.path, + } + .fail(); + } + if meta.content_length() != object.length { + return error::InvalidCopyParameterSnafu { + key: "object_length", + value: &object.path, + } + .fail(); + } + } + Ok(PreparedPackedImport { + backend, + index, + tables, + parallelism: crate::statement::database_copy::parse_parallelism_from_option_map( + &req.with, + ), + }) + } + + pub async fn import_packed( + &self, + plan: PreparedPackedImport, + cancellation: &CancellationToken, + ctx: QueryContextRef, + ) -> Result { + let windows = PackReadWindows::new(plan.backend.object_store.clone()); + let mut pending = PendingPackedInserts::new(plan.parallelism); + let objects: HashMap<_, _> = plan + .index + .objects + .iter() + .map(|o| (o.path.as_str(), o)) + .collect(); + let result: Result<()> = async { + for (entry, table) in plan.index.tables.iter().zip(plan.tables) { + if cancellation.is_cancelled() { + return error::PackedImportCancelledSnafu.fail(); + } + let object = objects.get(entry.object.as_str()).ok_or_else(|| { + error::InvalidCopyParameterSnafu { + key: "object", + value: &entry.object, + } + .build() + })?; + match object.kind { + ObjectKind::Pack if entry.length <= WINDOW_SIZE as u64 => { + let reader = PackedParquetReader::new( + Arc::clone(&windows), + object.path.clone(), + object.length, + entry.offset, + entry.length, + ) + .context(error::ReadParquetMetadataSnafu)?; + self.copy_indexed_parquet( + reader, + table, + entry.row_count, + &mut pending, + cancellation, + ctx.clone(), + ) + .await? + } + ObjectKind::Pack | ObjectKind::Parquet => { + // The virtual file keeps footer/column offsets relative to this + // stream. Parquet decoder allocations are separate from the cache. + let reader = plan + .backend + .object_store + .reader_with(&object.path) + .chunk(256 * 1024) + .await + .context(error::ReadObjectSnafu { path: &object.path })? + .into_futures_async_read(entry.offset..entry.offset + entry.length) + .await + .context(error::ReadObjectSnafu { path: &object.path })? + .compat(); + self.copy_indexed_parquet( + reader, + table, + entry.row_count, + &mut pending, + cancellation, + ctx.clone(), + ) + .await? + } + }; + } + Ok(()) + } + .await; + let drained = pending.drain().await; + common_telemetry::debug!( + peak_pending_bytes = pending.peak_pending_bytes, + peak_batch_bytes = pending.peak_batch_bytes, + "Packed import insertion payloads; codec and compressed windows accounted separately" + ); + result?; + drained?; + Ok(Output::new( + common_query::OutputData::AffectedRows(pending.rows), + common_query::OutputMeta::new_with_cost(pending.cost), + )) + } +} + +/// Retains all started writes until completion, including after the first error. +pub(crate) struct PendingPackedInserts { + tasks: FuturesUnordered)>>, + bytes: usize, + parallelism: usize, + rows: usize, + cost: usize, + peak_pending_bytes: usize, + peak_batch_bytes: usize, +} + +impl PendingPackedInserts { + fn new(parallelism: usize) -> Self { + Self { + tasks: FuturesUnordered::new(), + bytes: 0, + parallelism, + rows: 0, + cost: 0, + peak_pending_bytes: 0, + peak_batch_bytes: 0, + } + } + + /// Backpressure before decoding, while allowing one current batch outside the queue. + pub(crate) async fn before_decode(&mut self) -> Result<()> { + while let Some(Some(output)) = self.tasks.next().now_or_never() { + self.complete(output)?; + } + while self.tasks.len() >= self.parallelism || self.bytes >= 32 * 1024 * 1024 { + self.complete_one().await?; + } + Ok(()) + } + + pub(crate) async fn admit( + &mut self, + bytes: usize, + insert: impl std::future::Future> + Send + 'static, + cancellation: &CancellationToken, + ) -> Result<()> { + while let Some(Some(output)) = self.tasks.next().now_or_never() { + self.complete(output)?; + } + self.peak_batch_bytes = self.peak_batch_bytes.max(bytes); + const BUDGET: usize = 32 * 1024 * 1024; + if bytes > BUDGET { + self.drain().await?; + if cancellation.is_cancelled() { + return error::PackedImportCancelledSnafu.fail(); + } + let (rows, cost) = insert.await?.extract_rows_and_cost(); + self.rows += rows; + self.cost += cost; + return Ok(()); + } + while self.tasks.len() >= self.parallelism || self.bytes + bytes > BUDGET { + self.complete_one().await?; + } + if cancellation.is_cancelled() { + return error::PackedImportCancelledSnafu.fail(); + } + self.bytes += bytes; + self.peak_pending_bytes = self.peak_pending_bytes.max(self.bytes); + self.tasks.push(common_runtime::spawn_query( + async move { (bytes, insert.await) }, + )); + Ok(()) + } + + async fn complete_one(&mut self) -> Result<()> { + if let Some(output) = self.tasks.next().await { + self.complete(output)?; + } + Ok(()) + } + + fn complete( + &mut self, + output: std::result::Result<(usize, Result), common_runtime::JoinError>, + ) -> Result<()> { + let (bytes, output) = output.context(error::JoinTaskSnafu)?; + self.bytes -= bytes; + let (rows, cost) = output?.extract_rows_and_cost(); + self.rows += rows; + self.cost += cost; + Ok(()) + } + + async fn drain(&mut self) -> Result<()> { + let mut first = None; + while !self.tasks.is_empty() { + if let Err(error) = self.complete_one().await { + first.get_or_insert(error); + } + } + first.map_or(Ok(()), Err) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[tokio::test] + async fn oversized_batch_drains_and_runs_exclusively() { + let cancellation = CancellationToken::new(); + let mut queue = PendingPackedInserts::new(4); + let completed = Arc::new(AtomicUsize::new(0)); + for _ in 0..2 { + let completed = completed.clone(); + queue + .admit( + 1024, + async move { + completed.fetch_add(1, Ordering::SeqCst); + Ok(Output::new_with_affected_rows(1)) + }, + &cancellation, + ) + .await + .unwrap(); + } + let observed = completed.clone(); + queue + .admit( + 40 * 1024 * 1024, + async move { + assert_eq!(observed.load(Ordering::SeqCst), 2); + Ok(Output::new_with_affected_rows(3)) + }, + &cancellation, + ) + .await + .unwrap(); + assert_eq!(queue.rows, 5); + assert!(queue.tasks.is_empty()); + assert!(queue.peak_pending_bytes <= 32 * 1024 * 1024); + assert_eq!(queue.peak_batch_bytes, 40 * 1024 * 1024); + } + + #[tokio::test] + async fn concurrent_insert_queue_drains_after_error_and_cancel() { + let cancellation = CancellationToken::new(); + let mut queue = PendingPackedInserts::new(2); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let completed = Arc::new(AtomicUsize::new(0)); + for fail in [true, false] { + let barrier = barrier.clone(); + let completed = completed.clone(); + queue + .admit( + 16 * 1024 * 1024, + async move { + barrier.wait().await; + completed.fetch_add(1, Ordering::SeqCst); + if fail { + error::InvalidCopyParameterSnafu { + key: "test", + value: "failure", + } + .fail() + } else { + Ok(Output::new_with_affected_rows(1)) + } + }, + &cancellation, + ) + .await + .unwrap(); + } + cancellation.cancel(); + assert!( + queue + .admit(1, async { panic!("must not start") }, &cancellation) + .await + .is_err() + ); + let _ = queue.drain().await; + assert_eq!(completed.load(Ordering::SeqCst), 2); + assert_eq!(queue.bytes, 0); + assert!(queue.tasks.is_empty()); + } +} diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index fd82079c016..7f2688d30cf 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -1377,6 +1377,12 @@ impl HttpServer { fn route_sql(api_state: ApiState) -> Router { Router::new() + .route( + "/capabilities", + routing::get(|| async { + axum::Json(serde_json::json!({"metric_packed_import": 1})) + }), + ) .route( "/sql", routing::get(handler::sql).post(handler::sql).layer( @@ -1865,6 +1871,38 @@ mod test { assert!(!is_api_listener_path("/health")); } + #[tokio::test] + async fn packed_capability_requires_auth_on_full_and_api_listeners() { + let (tx, _rx) = mpsc::channel(1); + let provider = + auth::static_user_provider_from_option("static_user_provider:cmd:user=password") + .unwrap(); + let (full, api) = HttpServerBuilder::new(HttpOptions { + enable_api_server: true, + ..Default::default() + }) + .with_sql_handler(Arc::new(DummyInstance { _tx: tx })) + .with_user_provider(Arc::new(provider)) + .build_servers(); + for server in [full, api.unwrap()] { + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + assert_eq!( + client.get("/v1/capabilities").send().await.status(), + StatusCode::UNAUTHORIZED + ); + let response = client + .get("/v1/capabilities") + .header("Authorization", "Basic dXNlcjpwYXNzd29yZA==") + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.json::().await, + serde_json::json!({"metric_packed_import": 1}) + ); + } + } + #[tokio::test] pub async fn test_extra_router_respects_api_surface() { // Downstream routes added through `with_extra_router()` (e.g. enterprise diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index b5603bd29df..7a7d0cba72a 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -98,6 +98,7 @@ zstd.workspace = true [dev-dependencies] catalog = { workspace = true, features = ["testing"] } +cli.workspace = true client = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-procedure-test.workspace = true diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index e87bf387979..92888cced3c 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -747,6 +747,643 @@ async fn database_export_preserves_valid_table_names() { } } +#[tokio::test(flavor = "multi_thread")] +async fn packed_copy_standalone_heterogeneous_streams() { + use common_datasource::packed_snapshot::{ObjectKind, PackIndex, PackObject, PackTable}; + let standalone = GreptimeDbStandaloneBuilder::new("packed_copy") + .build() + .await; + let instance = standalone.fe_instance(); + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let data_directory = destination.path().join("data/public/1"); + std::fs::create_dir_all(&data_directory).unwrap(); + let directory = data_directory.as_path(); + let mut index = PackIndex { + version: 1, + objects: vec![], + tables: vec![], + }; + let mut pack = Vec::new(); + let mut expected = Vec::new(); + let mut creates = Vec::new(); + let mut physical_creates = Vec::new(); + let mut physical_ids = Vec::new(); + // Metric tag schemas vary; the standalone file has a different value type. + for (i, (name, typ, value)) in [ + ("literal.name", "DOUBLE", Some("42")), + ("strings", "DOUBLE", Some("43,'tag'")), + ("empty", "DOUBLE", None), + ("standalone", "STRING", Some("'standalone'")), + ] + .into_iter() + .enumerate() + { + let create = if i < 3 { + let physical = format!( + "CREATE TABLE IF NOT EXISTS physical_{i} (ts TIMESTAMP TIME INDEX, val {typ}) ENGINE=metric WITH(physical_metric_table='')" + ); + sql(instance, &physical).await; + physical_creates.push(physical); + physical_ids.push( + table(instance, &format!("physical_{i}")) + .await + .table_info() + .ident + .table_id, + ); + let tag = if i == 1 { + ", host STRING PRIMARY KEY" + } else { + "" + }; + format!( + "CREATE TABLE \"{name}\" (ts TIMESTAMP TIME INDEX, val {typ}{tag}) ENGINE=metric WITH(on_physical_table='physical_{i}')" + ) + } else { + format!("CREATE TABLE \"{name}\" (ts TIMESTAMP TIME INDEX, val {typ})") + }; + sql(instance, &create).await; + creates.push(create.clone()); + if let Some(value) = value { + let columns = if i == 1 { "ts,val,host" } else { "ts,val" }; + sql( + instance, + &format!("INSERT INTO \"{name}\" ({columns}) VALUES (1,{value})"), + ) + .await; + } + expected.push(values(instance, &format!("SELECT * FROM \"{name}\"")).await); + let path = directory.join(format!("table-{i}.parquet")); + sql( + instance, + &format!( + "COPY \"{name}\" TO '{}' WITH (FORMAT='parquet')", + path.display() + ), + ) + .await; + let bytes = std::fs::read(&path).unwrap(); + let (object, offset) = if i == 3 { + index.objects.push(PackObject { + path: format!("table-{i}.parquet"), + kind: ObjectKind::Parquet, + length: bytes.len() as u64, + }); + (format!("table-{i}.parquet"), 0) + } else { + let offset = pack.len() as u64; + pack.extend_from_slice(&bytes); + std::fs::remove_file(&path).unwrap(); + ("pack-0.bin".into(), offset) + }; + index.tables.push(PackTable { + table_name: name.into(), + object, + offset, + length: bytes.len() as u64, + row_count: u64::from(value.is_some()), + }); + sql(instance, &format!("DROP TABLE \"{name}\"")).await; + sql(instance, &create).await; + } + index.objects.push(PackObject { + path: "pack-0.bin".into(), + kind: ObjectKind::Pack, + length: pack.len() as u64, + }); + std::fs::write(directory.join("pack-0.bin"), pack).unwrap(); + std::fs::write( + directory.join("pack-index.json"), + serde_json::to_vec(&index).unwrap(), + ) + .unwrap(); + struct DenyPackedTable(std::sync::atomic::AtomicBool); + impl auth::PermissionChecker for DenyPackedTable { + fn check_permission( + &self, + _: auth::UserInfoRef, + _: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + fn check_permission_with_table_targets( + &self, + _: auth::UserInfoRef, + req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + let denied = self.0.load(std::sync::atomic::Ordering::SeqCst) + && !req.is_readonly() + && matches!(targets, auth::PermissionTableTargets::Resolved(ref tables) if tables.iter().any(|t| t.table == "strings")); + Ok(if denied { + auth::PermissionResp::Reject + } else { + auth::PermissionResp::Allow + }) + } + } + let checker = Arc::new(DenyPackedTable(std::sync::atomic::AtomicBool::new(true))); + instance + .plugins() + .insert::(checker.clone()); + let statement = format!( + "COPY DATABASE public FROM '{}/' WITH (FORMAT='parquet', metric_data_layout='packed', parallelism=2)", + directory.display() + ); + let denied = servers::query_handler::sql::SqlQueryHandler::do_query( + instance.as_ref(), + &statement, + QueryContext::arc(), + ) + .await; + assert!(denied.into_iter().all(|r| r.is_err())); + checker.0.store(false, std::sync::atomic::Ordering::SeqCst); + for entry in &index.tables { + assert!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)) + .await + .is_empty() + ); + } + let output = sql(instance, &format!("COPY DATABASE public FROM '{}/' WITH (FORMAT='parquet', metric_data_layout='packed', parallelism=2)", directory.display())).await; + assert!(matches!(output.data, OutputData::AffectedRows(3))); + for (entry, expected) in index.tables.iter().zip(&expected) { + assert_eq!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)).await, + *expected + ); + } + if let Ok(bucket) = std::env::var("GT_S3_BUCKET") + && !bucket.is_empty() + { + let prefix = format!( + "packed-reader/{}/", + destination.path().file_name().unwrap().to_str().unwrap() + ); + let mut connection = std::collections::HashMap::from([ + ( + "access_key_id".into(), + std::env::var("GT_S3_ACCESS_KEY_ID").unwrap(), + ), + ( + "secret_access_key".into(), + std::env::var("GT_S3_ACCESS_KEY").unwrap(), + ), + ("region".into(), std::env::var("GT_S3_REGION").unwrap()), + ]); + if let Ok(endpoint) = std::env::var("GT_S3_ENDPOINT_URL") { + connection.insert("endpoint".into(), endpoint); + } + let location = format!("s3://{bucket}/{prefix}"); + let store = common_datasource::object_store::build_backend( + &location, + &connection, + &common_datasource::object_store::LocalFileAccess::Disabled, + ) + .await + .unwrap(); + for path in index + .objects + .iter() + .map(|o| o.path.as_str()) + .chain(["pack-index.json"]) + { + store + .write(path, std::fs::read(directory.join(path)).unwrap()) + .await + .unwrap(); + } + for (entry, create) in index.tables.iter().zip(&creates) { + sql(instance, &format!("DROP TABLE \"{}\"", entry.table_name)).await; + sql(instance, create).await; + } + let connection_sql = connection + .iter() + .map(|(k, v)| format!("{k}='{}'", v.replace('\'', "''"))) + .collect::>() + .join(","); + sql(instance, &format!("COPY DATABASE public FROM '{location}' WITH (FORMAT='parquet', metric_data_layout='packed', parallelism=2) CONNECTION ({connection_sql})")).await; + for (entry, expected) in index.tables.iter().zip(&expected) { + assert_eq!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)).await, + *expected + ); + } + for path in index + .objects + .iter() + .map(|o| o.path.as_str()) + .chain(["pack-index.json"]) + { + store.delete(path).await.unwrap(); + } + } + // Run the real CLI importer through HTTP against the same generated fixture. + use clap::Parser; + use cli::export_v2::manifest::{ChunkMeta, DataFormat, Manifest, TimeRange}; + let mut manifest = Manifest::new_full( + "greptime".into(), + vec!["public".into()], + TimeRange::unbounded(), + DataFormat::Parquet, + ); + manifest.version = 2; + manifest.data_layout = Some("metric-parquet-packs".into()); + let mut chunk = ChunkMeta::new(1, TimeRange::unbounded()); + chunk.mark_completed( + index + .objects + .iter() + .map(|o| format!("data/public/1/{}", o.path)) + .chain(["data/public/1/pack-index.json".into()]) + .collect(), + None, + ); + manifest.chunks.push(chunk); + let root = destination.path(); + std::fs::create_dir_all(root.join("schema/ddl")).unwrap(); + std::fs::write( + root.join("manifest.json"), + serde_json::to_vec(&manifest).unwrap(), + ) + .unwrap(); + std::fs::write( + root.join("schema/ddl/public.sql"), + physical_creates + .iter() + .chain(&creates) + .cloned() + .collect::>() + .join(";\n") + + ";", + ) + .unwrap(); + for entry in &index.tables { + sql(instance, &format!("DROP TABLE \"{}\"", entry.table_name)).await; + } + for i in 0..physical_ids.len() { + sql(instance, &format!("DROP TABLE physical_{i}")).await; + } + let server = servers::http::HttpServerBuilder::new(Default::default()) + .with_sql_handler(instance.clone()) + .with_user_provider(Arc::new( + auth::static_user_provider_from_option("static_user_provider:cmd:user=password") + .unwrap(), + )) + .build(); + let client = + servers::http::test_helpers::TestClient::new(server.build(server.make_app()).unwrap()) + .await; + let state_path = root.join("import-state.json"); + let snapshot_uri = url::Url::from_file_path(root).unwrap(); + let command = cli::import_v2::ImportV2Command::parse_from([ + "import-v2", + "--addr", + client.base_url().trim_start_matches("http://"), + "--from", + snapshot_uri.as_str(), + "--state-path", + state_path.to_str().unwrap(), + "--auth-basic", + "user:password", + "--no-proxy", + "--progress", + "never", + ]); + let importer = command.build().await.unwrap(); + std::fs::write( + root.join("schema/schemas.json"), + serde_json::to_vec( + &serde_json::json!([{"catalog":"greptime","name":"public","options":{}}]), + ) + .unwrap(), + ) + .unwrap(); + #[derive(clap::Parser)] + struct VerifyArgs { + #[command(subcommand)] + command: cli::export_v2::ExportV2Command, + } + let verifier = + VerifyArgs::parse_from(["export-v2", "verify", "--snapshot", snapshot_uri.as_str()]) + .command + .build() + .await + .unwrap(); + verifier.do_work().await.unwrap(); + let original = serde_json::to_value(&manifest).unwrap(); + let initial_tables = values(instance, "SHOW TABLES").await; + let pack_path = directory.join("pack-0.bin"); + let pack_bytes = std::fs::read(&pack_path).unwrap(); + for case in [ + "schema_only", + "missing_chunks", + "missing_files", + "manifest_checksum", + "chunk_checksum", + "duplicate", + "extra", + "missing_object", + "truncated_object", + ] { + let mut invalid = original.clone(); + let expected_error = match case { + "schema_only" => { + invalid["schema_only"] = true.into(); + "schema_only" + } + "missing_chunks" => { + invalid.as_object_mut().unwrap().remove("chunks"); + "chunks array" + } + "missing_files" => { + invalid["chunks"][0] + .as_object_mut() + .unwrap() + .remove("files"); + "files array" + } + "manifest_checksum" => { + invalid["checksum"] = "unsupported".into(); + "checksums" + } + "chunk_checksum" => { + invalid["chunks"][0]["checksum"] = "unsupported".into(); + "checksums" + } + "duplicate" => { + let file = invalid["chunks"][0]["files"][0].clone(); + invalid["chunks"][0]["files"] + .as_array_mut() + .unwrap() + .push(file); + "duplicate" + } + "extra" => { + std::fs::write(directory.join("extra.bin"), b"extra").unwrap(); + invalid["chunks"][0]["files"] + .as_array_mut() + .unwrap() + .push("data/public/1/extra.bin".into()); + "inventory" + } + "missing_object" => { + std::fs::remove_file(&pack_path).unwrap(); + "expected length" + } + "truncated_object" => { + std::fs::write(&pack_path, &pack_bytes[..pack_bytes.len() - 1]).unwrap(); + "expected length" + } + _ => unreachable!(), + }; + std::fs::write( + root.join("manifest.json"), + serde_json::to_vec(&invalid).unwrap(), + ) + .unwrap(); + for resumed in [false, true] { + if resumed { + std::fs::write(&state_path, b"state must not be read or changed").unwrap(); + } + let error = importer.do_work().await.unwrap_err(); + assert!( + format!("{error:?}").contains(expected_error), + "{case}: {error:?}" + ); + assert_eq!( + values(instance, "SHOW TABLES").await, + initial_tables, + "{case}" + ); + if resumed { + assert_eq!( + std::fs::read(&state_path).unwrap(), + b"state must not be read or changed" + ); + std::fs::remove_file(&state_path).unwrap(); + } else { + assert!(!state_path.exists(), "{case}"); + } + } + assert!(verifier.do_work().await.is_err(), "{case}"); + std::fs::write(&pack_path, &pack_bytes).unwrap(); + if case == "extra" { + std::fs::remove_file(directory.join("extra.bin")).unwrap(); + } + } + std::fs::write( + root.join("manifest.json"), + serde_json::to_vec(&manifest).unwrap(), + ) + .unwrap(); + verifier.do_work().await.unwrap(); + importer.do_work().await.unwrap(); + assert!(!state_path.exists()); + for (i, old_id) in physical_ids.iter().enumerate() { + assert_ne!( + *old_id, + table(instance, &format!("physical_{i}")) + .await + .table_info() + .ident + .table_id + ); + } + for (entry, expected) in index.tables.iter().zip(&expected) { + assert_eq!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)).await, + *expected + ); + } + for entry in &index.tables { + sql(instance, &format!("DROP TABLE \"{}\"", entry.table_name)).await; + } + index.tables.last_mut().unwrap().row_count += 1; + std::fs::write( + directory.join("pack-index.json"), + serde_json::to_vec(&index).unwrap(), + ) + .unwrap(); + let failure = importer.do_work().await.unwrap_err(); + assert!( + format!("{failure:?}").contains("Chunk 1 import failed"), + "{failure:?}" + ); + let failed: serde_json::Value = + serde_json::from_slice(&std::fs::read(&state_path).unwrap()).unwrap(); + assert_eq!(failed["ddl_completed"], true); + assert_eq!(failed["tasks"][0]["status"], "failed"); + let saved_state = std::fs::read(&state_path).unwrap(); + std::fs::write(&pack_path, &pack_bytes[..pack_bytes.len() - 1]).unwrap(); + let preflight_error = importer.do_work().await.unwrap_err(); + assert!( + format!("{preflight_error:?}").contains("expected length"), + "{preflight_error:?}" + ); + assert_eq!(std::fs::read(&state_path).unwrap(), saved_state); + std::fs::write(&pack_path, &pack_bytes).unwrap(); + let failed_copy = servers::query_handler::sql::SqlQueryHandler::do_query( + instance.as_ref(), + &statement, + QueryContext::arc(), + ) + .await + .remove(0) + .err() + .unwrap(); + assert!( + format!("{failed_copy:?}").contains("row_count"), + "{failed_copy:?}" + ); + + // A failed chunk may have inserted rows; clear them before explicit replay. + for (entry, create) in index.tables.iter().zip(&creates) { + sql(instance, &format!("DROP TABLE \"{}\"", entry.table_name)).await; + sql(instance, create).await; + } + index.tables.last_mut().unwrap().row_count -= 1; + std::fs::write( + directory.join("pack-index.json"), + serde_json::to_vec(&index).unwrap(), + ) + .unwrap(); + importer.do_work().await.unwrap(); + assert!(!state_path.exists()); + for (entry, expected) in index.tables.iter().zip(&expected) { + assert_eq!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)).await, + *expected + ); + } + + // The same importer still restores the original per-table manifest layout. + let packed_bytes = std::fs::read(directory.join("pack-0.bin")).unwrap(); + let mut files = Vec::new(); + for entry in &index.tables { + let name = format!("{}.parquet", entry.table_name); + if entry.object == "pack-0.bin" { + std::fs::write( + directory.join(&name), + &packed_bytes[entry.offset as usize..(entry.offset + entry.length) as usize], + ) + .unwrap(); + } else { + std::fs::rename(directory.join(&entry.object), directory.join(&name)).unwrap(); + } + files.push(format!("data/public/1/{name}")); + sql(instance, &format!("DROP TABLE \"{}\"", entry.table_name)).await; + } + manifest.version = 1; + manifest.data_layout = None; + manifest.chunks[0].files = files; + std::fs::write( + root.join("manifest.json"), + serde_json::to_vec(&manifest).unwrap(), + ) + .unwrap(); + importer.do_work().await.unwrap(); + for (entry, expected) in index.tables.iter().zip(&expected) { + assert_eq!( + values(instance, &format!("SELECT * FROM \"{}\"", entry.table_name)).await, + *expected + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn packed_copy_large_column_in_middle_stream() { + use common_datasource::packed_snapshot::{ObjectKind, PackIndex, PackObject, PackTable}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::basic::Compression; + use datafusion::parquet::file::properties::WriterProperties; + use datatypes::arrow::array::{ + ArrayRef, BooleanArray, Int64Array, StringArray, TimestampMillisecondArray, + }; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + let standalone = GreptimeDbStandaloneBuilder::new("packed_large_column") + .build() + .await; + let instance = standalone.fe_instance(); + let directory = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let large = "x".repeat(20 * 1024 * 1024); + let arrays: [ArrayRef; 3] = [ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![large.as_str()])), + Arc::new(BooleanArray::from(vec![true])), + ]; + let mut packed = Vec::new(); + let mut index = PackIndex { + version: 1, + objects: vec![], + tables: vec![], + }; + for (i, (array, typ)) in arrays + .into_iter() + .zip(["BIGINT", "STRING", "BOOLEAN"]) + .enumerate() + { + sql( + instance, + &format!("CREATE TABLE restored_{i} (ts TIMESTAMP TIME INDEX, val {typ})"), + ) + .await; + let timestamps: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1])); + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamps.data_type().clone(), false), + Field::new("val", array.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![timestamps, array]).unwrap(); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .build(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + if i == 1 { + assert!(metadata.row_group(0).column(1).compressed_size() > 16 * 1024 * 1024); + assert!(!packed.is_empty()); + } + index.tables.push(PackTable { + table_name: format!("restored_{i}"), + object: "pack-0.bin".into(), + offset: packed.len() as u64, + length: bytes.len() as u64, + row_count: 1, + }); + packed.extend(bytes); + } + index.objects.push(PackObject { + path: "pack-0.bin".into(), + kind: ObjectKind::Pack, + length: packed.len() as u64, + }); + std::fs::write(directory.path().join("pack-0.bin"), packed).unwrap(); + std::fs::write( + directory.path().join("pack-index.json"), + serde_json::to_vec(&index).unwrap(), + ) + .unwrap(); + let output = sql(instance, &format!("COPY DATABASE public FROM '{}/' WITH (FORMAT='parquet', metric_data_layout='packed', parallelism=2)", directory.path().display())).await; + assert!(matches!(output.data, OutputData::AffectedRows(3))); + for (i, expected) in [ + datatypes::value::Value::Int64(42), + datatypes::value::Value::String(large.into()), + datatypes::value::Value::Boolean(true), + ] + .into_iter() + .enumerate() + { + assert_eq!( + values(instance, &format!("SELECT val FROM restored_{i}")).await, + vec![vec![expected]] + ); + } +} + #[derive(clap::Parser)] struct ExportDataCli { #[command(subcommand)] @@ -884,6 +1521,30 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { } let (addr, server) = export_http(instance.clone()).await; let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + for experimental in [false, true] { + for layout in ["packed", "invalid"] { + let path = destination + .path() + .join(format!("unsupported-{experimental}-{layout}")); + let uri = url::Url::from_directory_path(&path).unwrap(); + let statement = format!( + "COPY DATABASE public TO '{uri}' WITH (FORMAT='parquet', experimental_metric_export='{experimental}', metric_data_layout='{layout}')" + ); + let result = servers::query_handler::sql::SqlQueryHandler::do_query( + instance.as_ref(), + &statement, + QueryContext::arc(), + ) + .await + .remove(0); + let error = result.expect_err("export must reject import-only layouts"); + assert!( + format!("{error:?}").contains("metric_data_layout"), + "{error:?}" + ); + assert!(!path.exists()); + } + } let (uri, store, storage_args) = if s3 { let endpoint = std::env::var("GT_S3_ENDPOINT_URL").unwrap(); let bucket = std::env::var("GT_S3_BUCKET").unwrap(); diff --git a/tests/compatibility/cases/snapshot_parquet_copy/case.toml b/tests/compatibility/cases/snapshot_parquet_copy/case.toml new file mode 100644 index 00000000000..28b6a4b69da --- /dev/null +++ b/tests/compatibility/cases/snapshot_parquet_copy/case.toml @@ -0,0 +1,8 @@ +name = "snapshot_parquet_copy" +reason = "Preserve explicit-file Parquet restore from snapshots written before packed import support." +introduced_by = "packed snapshot reader" +topologies = ["standalone"] +from_range = ["*"] +to_range = ["*"] +features = ["copy", "parquet"] +owner = "GreptimeTeam" diff --git a/tests/compatibility/cases/snapshot_parquet_copy/setup.sql b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql new file mode 100644 index 00000000000..8342e1847b5 --- /dev/null +++ b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql @@ -0,0 +1,4 @@ +CREATE TABLE snapshot_values (ts TIMESTAMP TIME INDEX, val BIGINT); +INSERT INTO snapshot_values VALUES (1, 42), (2, 7); +COPY snapshot_values TO '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); +TRUNCATE TABLE snapshot_values; diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.result b/tests/compatibility/cases/snapshot_parquet_copy/verify.result new file mode 100644 index 00000000000..d3e4d2483ab --- /dev/null +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.result @@ -0,0 +1,12 @@ +COPY snapshot_values FROM '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); + +Affected Rows: 2 + +SELECT * FROM snapshot_values ORDER BY ts; + ++-------------------------+-----+ +| ts | val | ++-------------------------+-----+ +| 1970-01-01T00:00:00.001 | 42 | +| 1970-01-01T00:00:00.002 | 7 | ++-------------------------+-----+ diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.sql b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql new file mode 100644 index 00000000000..c782dcacb1b --- /dev/null +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql @@ -0,0 +1,2 @@ +COPY snapshot_values FROM '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); +SELECT * FROM snapshot_values ORDER BY ts; From 4a980f350f927c3becf07a228adb98951999bf3f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 00:13:13 +0000 Subject: [PATCH 44/63] docs: make packed snapshots part of the Metric export/import RFC (#9247) * docs: revise Metric snapshots around packed Parquet objects Signed-off-by: jeremyhi * docs: summarize experiment conclusions and bound multipart fallback Signed-off-by: jeremyhi * docs: specify packed snapshot persistence contract Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- docs/rfcs/2026-09-11-metric-export-import.md | 473 ++++++++++++++---- .../experiment-results.md | 135 ++--- 2 files changed, 414 insertions(+), 194 deletions(-) diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 1a984211e4d..3d2c93e4695 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -2,21 +2,27 @@ Feature Name: Metric Export and Import Optimization Tracking Issue: https://github.com/GreptimeTeam/greptimedb/issues/9120 Date: 2026-09-11 +Updated: 2026-09-21 Status: Draft for discussion --- # Summary -Group the selected Metric logical tables by physical table. Query each physical -table across all its regions, order the rows by `__table_id`, and project each -logical schema into the existing logical-table Parquet files. Restore the same -V2 snapshot with batched logical-table DDL and the existing data COPY path. A -separate generic COPY fix removes repeated directory scans. +Scan Metric tables by physical group, project each logical schema, and pack +small, complete logical-table Parquet streams into shared objects. An index maps +each table to its object and byte range. Restore through a bounded range reader +and normal target-table insertion, with batched logical-table DDL. -The design preserves `manifest.version = 1` and the current V2 layout. It does -not introduce physical backup files or require importing source table IDs. -Merged data writes remain a first-release scope decision, pending representative -measurements after the COPY fix. +Packing is required for the first release of this optimization. Packed snapshots +use `manifest.version = 2`; new importers read versions 1 and 2, while old +importers reject version 2 before executing DDL. The command names remain +`export-v2` and `import-v2`. Existing version-1 snapshots retain their layout. + +Each storage write and multipart data part is at most 8 MiB. Objects can be +larger, and a logical table is not split to meet that write size. Shared physical +scans, bounded writer concurrency, removal of redundant object checks, schema +export batching and batch CREATE together address the complete transfer cost. +Merged database insert requests remain a separate follow-up. # Motivation and evidence @@ -25,11 +31,13 @@ includes reports of exports taking days for roughly 100,000 logical tables with only several GiB of data. Table count, physical schema width and data volume must be measured separately; total database bytes do not describe the transfer cost. -Local experiments support three independent optimizations: shared physical scans, -batched logical CREATE, and removal of repeated file lookup. The -[experiment results](2026-09-11-metric-export-import/experiment-results.md) -contain the measurements, controls and coverage limits. The COPY improvement -uses per-table writes; merged-write benefits remain unmeasured. +The [experiment conclusions](2026-09-11-metric-export-import/experiment-results.md) +support shared scans, batch DDL, explicit-file COPY and packed-object transfer. +Packing with a bounded range reader substantially reduces storage requests and +improves full export and restore for many small logical tables. Benefits depend +on request latency, and memory use increases. Remaining CREATE work makes batch +DDL part of the first release. These local experiments establish direction; +production V2 recovery and release performance still require acceptance tests. # Terminology @@ -41,10 +49,12 @@ uses per-table writes; merged-write benefits remain unmeasured. | Export unit | The selected logical tables sharing one physical table, within one schema and V2 time chunk. This is an export grouping, not a storage object or completion checkpoint. | | Query execution partition | A stream of work in the query plan. Frontend `target_partitions` influences the partition count; execution partitions are distinct from storage regions and export units. | | Logical-table descriptor | Captured catalog metadata: source ID, name, physical association, column names/types/order, time index and primary key. Capturing and reusing it provides no snapshot or locking guarantee. | -| Parquet writer | The writer for one logical table's output file in an export unit. | +| Parquet writer | An encoder for one complete logical-table Parquet stream, including its schema and footer. | +| Pack | An object containing consecutive complete Parquet streams, confined to one schema and time chunk. | +| Pack index | The table-to-object/range mapping and object lengths for one schema and time chunk, including standalone objects. | | Merged writes | Combining insert requests for multiple target logical tables into a region batch dispatch while retaining each table's identity, schema and routing. | -# Physical export and logical files +# Physical export and logical Parquet streams ## Query and routing @@ -77,14 +87,18 @@ physical table directly; it does not first write and reread `phy.parquet`. ```text physical table query: selected column union + __table_id, ordered by ID → reject rows whose IDs are absent from the captured descriptor set - → contiguous rows for cpu → project (ts, host, cpu) → cpu.parquet - → contiguous rows for requests → project (ts, host, service, requests) → requests.parquet + → contiguous cpu rows → project (ts, host, cpu) → cpu Parquet + → contiguous requests rows → project (ts, host, service, requests) → requests Parquet + → append complete small Parquet streams to a shared pack and record their ranges ``` -Only one Parquet writer is active per export unit. At an ID transition, close the -previous writer and open the next. Verify non-null UInt32 IDs and monotonic order -across input batches. Create a valid zero-row file for every selected empty table. -Never put `__table_id`, `__tsid`, or another table's columns into logical files. +The router sends each table's contiguous batches to one encoder. At an ID +transition, finish that encoder's input; earlier encoders may finish concurrently +under a request-wide writer budget. Append only complete Parquet streams, so +completion order need not match source table-ID order. Verify non-null UInt32 IDs +and monotonic input order across batches. Encode a valid zero-row Parquet for +every selected empty table. Never include `__table_id`, `__tsid`, or another +table's columns in a logical Parquet stream. The captured catalog IDs are the routing whitelist; physical rows are not the source of table membership. Dropped tables can leave residual rows. Predicate @@ -109,6 +123,36 @@ with query memory accounting, bounded conversion slices, row-group limits and cancellation. Reject an oversized row explicitly. These controls do not provide a hard process RSS guarantee. +Start with a 64 MiB request-wide budget for retained/queued Arrow payload and an +8 MiB encoded small-table buffer per admitted writer. When a stream outgrows that +buffer, flush its prefix into a standalone object and continue streaming the +same Parquet file. Bound encoder and upload concurrency by the request's +parallelism, shared across physical groups and ordinary-table work. Count shared +Arrow backing allocations and buffers retained by in-flight tasks, not only the +visible slice sizes. Codec and query reservations remain separate controls. + +Finishers retain their writer permits until their buffers are appended or +released. This prevents completed small files from forming an unbounded queue. +Each upload submission is at most 8 MiB, including standalone files and indexes. +CLI chunk concurrency multiplies per-request budgets and must be included in +resource sizing. Index and catalog memory scale with table count and are measured +separately. These initial budgets require the release resource tests below. + +## Schema export and destination checks + +Reuse an HTTP client and retrieve SHOW CREATE results in bounded multi-statement +requests, initially at most 128 statements and 128 KiB of SQL. Preserve statement +order, identifier quoting, per-statement errors and authorization. Respect response +limits and reject an oversized single request rather than bypassing the byte cap. +This reduces round trips without changing database/physical/logical/view DDL order. + +Reuse destination validation only within the same export request and exclusively +owned chunk directory. Check the destination before admission, and preserve +existing collision handling and local-file access rules for independent COPY +operations. An environment switch that globally skips HEAD is not the production +interface. A preflight check alone does not provide ownership against a live +overlapping exporter. + ## Ordering and multiple regions The first release covers every selected Metric physical table, any region count, @@ -134,21 +178,195 @@ export to per-logical-table queries. The # Snapshot and metadata contract -Retain the current implementation's layout: +## Versions and capability checks -```text -manifest.json # version 1 -schema/schemas.json # schema index -schema/ddl/public.sql # canonical database/table/view DDL -data/public/1/cpu.parquet -data/public/1/requests.parquet +Packed Metric exports write manifest version 2 with a required data-layout marker +`metric-parquet-packs`. Parquet remains the encoding of each table. Ordinary +exports can keep version 1. A version-2 snapshot indexes all data in its schema +chunks, including ordinary tables stored as standalone Parquet objects. + +Before creating or replacing a snapshot, the CLI checks the source server's +explicit packed-export capability. Before target DDL or data writes, the importer +checks the manifest version/layout and the target's packed-import capability. +An unsupported packed layout fails; it must not fall through to listing `.parquet` +files and silently omit packed tables. Batch-DDL capability is checked separately. + +| Reader and snapshot | Behavior | +| --- | --- | +| New importer, version 1 | Use the existing file layout and normal restore semantics. | +| New importer, version 2 | Require the packed layout and a target supporting packed import. | +| Old importer, version 2 | Reject the unsupported manifest version before DDL or data writes. | +| Unknown version or layout | Reject explicitly. | + +Update create, resume, verify, status and import together. Resume retains the +snapshot's original version and layout; it does not migrate an unfinished +version-1 snapshot. Validate old-reader rejection with supported release binaries. + +## Manifest and chunk schema + +The following is a complete version-2 manifest example. Field names and enum +values are part of the persisted contract; the timestamps and object names are +illustrative. + +```json +{ + "version": 2, + "data_layout": "metric-parquet-packs", + "snapshot_id": "123e4567-e89b-42d3-a456-426614174000", + "catalog": "greptime", + "schemas": ["public"], + "time_range": {"start": "2026-09-01T00:00:00Z", "end": "2026-09-02T00:00:00Z"}, + "schema_only": false, + "format": "parquet", + "chunks": [{ + "id": 1, + "time_range": {"start": "2026-09-01T00:00:00Z", "end": "2026-09-02T00:00:00Z"}, + "status": "completed", + "files": [ + "data/public/1/pack-000000.bin", + "data/public/1/table-000000.parquet", + "data/public/1/pack-index.json" + ] + }], + "created_at": "2026-09-19T00:00:00Z", + "updated_at": "2026-09-19T00:01:00Z" +} ``` -Schema path segments retain existing V2 encoding. Logical data filenames retain -COPY's raw table-name semantics and import's `file_stem` matching, including dots. -There is no existing general table-filename escaping contract to reuse. Reject -path separators for the fast path; changing that contract requires a separate -compatibility design. +All top-level and chunk fields shown are required in version 2. `snapshot_id` +is a UUID; timestamps are RFC3339; `time_range` retains the existing half-open interval semantics, +with an absent or null bound meaning unbounded on that side. Schema names are +literal identifiers and must be unique. Chunk IDs are unique positive `u32` +integers. Chunk statuses remain `pending`, `in_progress`, `completed`, `skipped` +and `failed`; a failed chunk may carry an optional string `error`. +`schema_only = true` requires an empty `chunks` array and no data objects; +schema-only exports continue to produce version 1. + +For version 1, an absent or null `data_layout` selects the existing layout and +all existing field defaults remain valid. Version 2 requires the exact non-null +`data_layout` above and `format = "parquet"`. Reject a layout marker on version +1, a missing or unknown version-2 layout, an unsupported version, or an +incompatible format before DDL, data writes or resume cleanup. Unknown manifest +and chunk fields may be ignored as optional metadata; a required decoding or +integrity feature must change the manifest version, not rely on such a field. + +For a completed version-2 chunk, `files` is the duplicate-free, order-independent +inventory of **all** schema indexes, pack objects and standalone objects in that +time chunk. It contains snapshot-relative paths using `/` on every platform. +For each schema, its subset must equal `pack-index.json` plus every object named +by that index, under `data///`. No other files belong +in this list, and DDL stays under `schema/`. Require canonical paths with no +leading slash, backslash or traversal; use the existing schema path encoding. +A filename extension never selects the layout or identifies a table. + +Every completed schema chunk has an index, including an empty index when its +DDL selects no data-bearing tables. A selected table with no rows still has a +zero-row Parquet entry. A chunk may be `skipped` with `files = []` only when +none of its schemas selects a data-bearing table, not merely when a time window +contains no rows. Pending, in-progress and failed chunks are never importable; +their files or index presence cannot establish completion. + +## Objects and index + +```text +manifest.json # version 2, metric-parquet-packs layout +schema/schemas.json +schema/ddl/public.sql +data/public/1/pack-000000.bin +data/public/1/table-000000.parquet # large logical or ordinary table +data/public/1/pack-index.json +``` + +Schema and chunk paths retain V2 encoding and final-directory placement. Object +names are generated within the chunk; table identity comes from the index, not +from a filename or a source ID. A dotted table name remains a literal identifier. +Version-1 file naming and `file_stem` matching remain unchanged. + +The matching `pack-index.json` has this shape; the lengths and row counts below +illustrate the mapping rather than an actual encoded fixture: + +```json +{ + "version": 1, + "objects": [ + {"path": "pack-000000.bin", "kind": "pack", "length": 4096}, + {"path": "table-000000.parquet", "kind": "parquet", "length": 8192} + ], + "tables": [ + {"table_name": "cpu", "object": "pack-000000.bin", "offset": 0, "length": 1536, "row_count": 10}, + {"table_name": "requests", "object": "pack-000000.bin", "offset": 1536, "length": 2560, "row_count": 20}, + {"table_name": "events", "object": "table-000000.parquet", "offset": 0, "length": 8192, "row_count": 100} + ] +} +``` + +The index version is independent of the manifest version and starts at 1. +All shown fields are required. Reject unknown index versions, fields or object +kinds. `length`, `offset` and `row_count` are unsigned `u64` integers; offsets +are zero-based bytes. Object paths are direct chunk children matching +`pack-[0-9]+.bin` or `table-[0-9]+.parquet` according to kind. The `object` field +references an object by its exact path. Table names are literal identifiers. +A standalone entry covers its entire object. Empty tables still have a nonempty, +valid zero-row Parquet stream. Each selected table has exactly one entry per +schema chunk; an empty index has `objects = []` and `tables = []`. + +Validate unique object paths and table names, object kinds/references, unsigned +range arithmetic, bounds, and nonoverlapping complete Parquet ranges. For a pack, +the ranges cover its bytes consecutively; standalone objects have one full-range +entry. Restrict object paths to generated direct children of the chunk directory; +reject absolute paths, separators and traversal. Resolve destination tables only +within the request's authorized catalog/schema. An index is not an authorization +grant. Match index membership to the snapshot table DDL, excluding Metric +physical tables and views, so a missing entry cannot silently omit a table. +Validate the index structure and membership before admitting writes, each stream's +Parquet schema through normal COPY checks, and its decoded row count against +the index before reporting that table complete. + +Use one streaming pack writer per schema chunk initially. Multiple physical +groups can feed it, with bounded multipart upload concurrency. Append each small +Parquet intact. A pack may exceed 8 MiB; 8 MiB limits one write/part, not the object. +Roll to a new pack at a table boundary before exhausting a backend object/part +limit. A standalone table is subject to both the backend object-size limit and +its multipart part-count limit. With at most N parts and an 8 MiB part cap, its +effective size ceiling is no greater than min(object-size limit, N × 8 MiB), and +can be lower if emitted parts are smaller. Before submitting an out-of-limit +part or exceeding the object limit, fail explicitly, abort the upload and leave +the chunk incomplete. Do not enlarge parts or split the logical table to bypass +these limits. Exercise this boundary with a reduced-part-limit test backend. + +Record object lengths while writing and include all standalone outputs in the +index. Restore need not list the directory or rediscover a pack's size for each +table. Stream index serialization in bounded writes; account for retained table +and index metadata independently of data buffers. + +## Validation and checksum coverage + +Before target DDL, import validates the manifest contract and, for every selected +schema and completed chunk, loads the required index, matches its inventory to +`files`, validates ranges and DDL membership, and checks object existence and +length once per object. Perform the same preflight on import resume, even when +local state records DDL or data tasks as completed. A malformed snapshot must +not reach the DDL executor or create new completed-task records. + +Version 2 retains the existing optional `checksum` fields at manifest and chunk +level, but writers omit both. Readers accept absent or null values and reject +non-null values as unsupported before DDL or resume cleanup. The existing +version-1 exporter does not compute these fields and `verify` does not validate +cryptographic checksums; its behavior remains unchanged. This version-2 contract +therefore provides no checksum coverage for the manifest, DDL, indexes or data +objects. Object-store ETags are not substituted for content checksums. Adding +checksums requires a versioned definition of the algorithm, covered bytes and +aggregation order before readers may claim to verify them. + +For version 2, `verify` applies the same metadata/index checks across all schemas, +checks completion and object lengths, and reports missing or unexpected data +files. Its result establishes structural completeness, not content-hash or +full row validation. Import additionally validates each Parquet stream and its +decoded row count through COPY; a same-length payload change can escape the +structural check. Detecting corrupt Parquet data during restore may leave partial +target writes, as described in the retry contract below. + +## DDL and target metadata Keep database, physical-table, logical/ordinary-table and view DDL in dependency order. Physical DDL describes the table and its options; logical DDL describes @@ -174,75 +392,109 @@ request size and the importer's DDL buffering separately. The production CLI needs an authenticated server operation that checks CREATE authorization for every member and reuses the ordinary DDL path. Transport, capability discovery and request limits must be agreed before accepting this -interface. An older server uses ordinary DDL; do not retry an ambiguously -completed batch through another path. +interface. Initially bound a batch to 128 statements and 1 MiB of SQL, with an +explicit oversized-statement error; separately bound the importer's parsed DDL +buffer. A target without batch-DDL capability uses ordinary DDL if it supports +the snapshot's data layout. Do not retry an ambiguously completed batch through +another path. -## Data COPY and generic directory fix +## Packed data COPY -Read each logical Parquet file through the current COPY path and write through -normal target logical-table routing. No source-ID remapping belongs in these -files, and batch DDL does not imply batched data insertion. +The CLI retains one `COPY DATABASE FROM` operation per `(chunk, schema)` task. +Add an explicit `metric_data_layout = 'packed'` COPY option, paired with +`FORMAT = 'parquet'`, for the advertised capability. The server validates the +index and selected destination tables, then schedules table streams by object +and offset. Use the same explicit layout option for packed export. The existing +Metric feature gate continues to control the experimental export path. -The generic fix changes explicit-file lookup from `stat → list parent → find` -to `stat → use path`, with a sandboxed file-type check for explicit local -inputs to preserve symlink filtering. Directory COPY retains listing, pattern -matching and file-type filtering. Backend authorization and error handling -remain in place. With N files in one directory, this removes roughly N² -directory-entry work. -COPY bypasses the common Lister's filename branch; the branch itself is unchanged. -This fix can ship independently of the Metric design. +A reader owned by that request translates each Parquet-relative read into a +bounded object range. Start with two shared 8 MiB read windows per request; +account for windows pinned by decoders before fetching replacements. Bound active +decoders and decoded batches separately. Large standalone tables stream through +the existing reader rather than being loaded into a small-file buffer. + +Reuse adjacent bytes for multiple tables and coalesce duplicate window fetches. +The cache is scoped to this request, authorized storage backend and snapshot; +it is not a process-global cache keyed only by object name. Use index lengths +and perform any necessary existence/size checks once per object, not per table. +Check short reads and Parquet schema/metadata errors normally. Objects must stay +stable for the duration of restore. + +Reuse COPY's schema validation, Parquet decoding and normal target-table insertion. +Preserve per-table authorization for all admitted entries and backend/local-file +access controls. Keep current schema selection. Per-table selection is a separate +future interface. Pack order changes storage access, not target table routing. + +## Version-1 COPY + +Keep the explicit-file optimization from #9126: `stat → use path`, including +the sandboxed local file-type check. Directory COPY retains listing, matching +and filtering. New importers continue to use this path for version-1 snapshots. +General removal of redundant restore metadata requests benefits that path too +and must be included in the baseline when measuring packing's incremental gain. ## Completion, interruption and retry -Keep V2 completion units: export chunk completion and import `(chunk, schema)` -data tasks, with the existing DDL-completed flag. Finish and close every required -file before recording completion. A closed subset of files is not a completed -chunk. On export retry, replace only the unfinished attempt's owned outputs; -never touch completed chunks. Cancellation must close/abort writers and leave -the chunk incomplete. +Keep export chunk completion, import `(chunk, schema)` data tasks and the existing +DDL-completed flag. Write data directly to final chunk paths. Close every pack +and standalone object, then finish the schema chunk's index. Only after all +required schema outputs and indexes succeed may V2 mark the chunk Completed. +Index presence alone does not establish completion; a failed manifest update +leaves the chunk incomplete. -Batch DDL is not a transaction across the restore. A later failure can leave -earlier batches created. Keep DDL completion false until all statements succeed, -then use the current durable state update. Data retries retain existing COPY -semantics and may replay an incomplete task; there is no new exactly-once -guarantee or per-file checkpoint in this proposal. +On failure, stop admission, cancel and drain started work, and abort unfinished +multipart uploads. Explicit retry cleans only owned outputs of an incomplete +chunk, including pack/index files, and reruns that chunk. Preserve completed +chunks and unrelated objects. A transport timeout does not establish that the +server stopped: require the existing stable-source and stopped-writer conditions +before cleanup/retry. Handle normal cancellation's multipart abort separately +from abandoned multipart uploads after a process crash; listing final objects +does not find the latter. + +The importer records a data task complete only after every indexed table succeeds. +Disable continue-on-error for this path: skipped failed tables must not produce a +successful task result. Partial target writes may remain after failure. Retry +retains existing COPY replay semantics, without per-table checkpoints or an +exactly-once guarantee. Cancel and drain active readers/inserts before returning +failure, so a later attempt cannot overlap leftover work from that request. + +Batch DDL is not transactional across restore. Earlier batches can remain after +a later failure. Keep DDL completion false until all statements succeed, then use +the existing durable state update. # Rationale and alternatives | Choice | Reason and trade-off | | --- | --- | -| Keep logical-table Parquet files | Reuses V2 DDL, file lookup and normal target CREATE/routing with fresh IDs. Physical backup files would need a new restore contract for mixed schemas and source metadata. The retained format still has one file per logical table per time chunk. | -| Scan by physical table | Shares query and scan work among logical tables. Per-logical-table COPY is simpler but retains the table-count overhead; physical scans instead pay for the selected column union and ID routing. | -| Global ID ordering and one Parquet writer per export unit | Contiguous IDs allow a file to be closed at each table transition. Unordered routing would require many open writers or intermediate buffering/files. Ordering bounds writer state, but query sorting and merging need separate resource controls. | -| Use query-engine ordering across regions | Reuses distributed planning and SQL semantics. An exporter-owned region scheduler would add merge, cancellation and retry coordination. The chosen approach must handle both streaming merge and additional-sort plans. | -| Batch DDL while retaining per-table data COPY | Reuses logical CREATE procedures and current data-task replay semantics. Merged writes may reduce dispatch/WAL overhead, but introduce partial-success and checkpoint questions and require a separate benefit comparison. | +| Pack complete logical Parquet streams | Reduces per-object requests while preserving different logical schemas, Parquet metadata and normal insertion. Requires a new snapshot version, index and reader. | +| Scan by physical table | Shares query/scan work; the selected column union and global ordering still need resource controls. | +| Ordered routing with bounded concurrent encoders | Keeps each table contiguous while overlapping encoding and output. Shared byte and task budgets bound downstream work. | +| 8 MiB write/part bound | Limits submission size and in-flight buffering while allowing larger objects. It is an initial balance, not a measured universal optimum. | +| Request-owned range windows | Shares object reads across tables without cross-request or cross-credential cache reuse. Decoder pinning must remain inside the byte budget. | +| Batch DDL and ordinary target insertion | Addresses remaining CREATE overhead while retaining target routing and existing replay semantics. Database merged writes remain a separate follow-up. | -# Unresolved questions +# Decisions still requiring implementation review -- Which authenticated batch-DDL operation and capability-discovery mechanism - should the CLI use? Agree request byte/table limits and DDL buffering bounds. -- What are the query memory, spill-space, writer and export-unit concurrency - limits? Agree the test workloads and resource/performance thresholds before - running the ordering acceptance benchmark. -- Which released import-v2 versions and object-store backends are compatibility - targets? Test those readers against the retained V2 format. -- What name and defaults should the `experimental_` export option use? -- Should merged writes ship in the first release? Compare a bounded merged-write - PoC against corrected per-table COPY with the same batch DDL and concurrency. - Bound decoded rows/bytes and in-flight requests, preserving normal logical - validation and target routing through the existing region batch dispatch. - Cover small/larger files, wide schemas and distributed/multi-region targets; - report repeated total/data times, rows/s, bytes/s, request counts and memory. - Agree a minimum useful gain and memory ceiling before the comparison. Validate - partial success, ambiguous completion and checkpoint advancement, then record - an explicit include/defer decision. +- Bind the authenticated batch-DDL operation to its production transport and + capability response before PR05; preserve ordinary CREATE validation, + per-table authorization and the batch/byte bounds above. +- Validate the initial writer and reader budgets against codec/query memory, + wide physical schemas, concurrent chunks and distributed sorting. Agree + deployment-level memory/spill ceilings before performance runs. +- Pin supported old-reader binaries for version-2 rejection and version-1 + compatibility tests. The first-release storage targets are local files and + S3-compatible storage exercised through MinIO. +- Decide when release acceptance permits removing the existing experimental + gate. Packed export/restore capability is required while mixed versions exist. # Integration and release acceptance Use existing COPY DATABASE orchestration to form Metric export units. Retain ordinary COPY for non-Metric tables and other formats, and direct COPY TABLE behavior. Gate the unfinished Metric fast path behind an `experimental_` -configuration option; disabling it selects the existing path. +configuration option; disabling it selects the existing path. Add packed-import +support before enabling packed snapshot creation. The writer and reader form one +first-release feature, even though they land in separate PRs. The [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120) owns the PR breakdown and implementation dependencies. Each implementation slice @@ -252,9 +504,52 @@ must verify its own correctness and failure behavior. Release acceptance require | --- | --- | | Export | Actual CLI export with authorization, multiple physical tables and regions, dense/sparse keys, empty/residual tables, supported Metric Parquet types, and logical schema/value equivalence. | | Restore | Authenticated remote batch DDL, per-table authorization, fresh target IDs and reconstructed metadata, dependency ordering, older-server behavior and batch failure/retry. | -| V2 lifecycle | Supported released readers; local and object-store interruption, owned-output cleanup, durable completion state and import replay. | +| V2 lifecycle | New-reader version-1/2 coverage, old-reader version-2 rejection, target capability preflight; local/object-store interruption, multipart abort, owned-output cleanup, index/manifest failure and import replay. | | Scale | Repeated release-build measurements varying table count up to approximately 100,000, rows/file, physical union width, regions and backend. Report median/spread, control execution order/cache, and include schema export and metadata discovery in end-to-end timing. | -| Open decisions | Agreed batch-DDL interface, resource limits, compatibility targets and experimental option; documented merged-write include/defer decision. | +| Open decisions | Agreed batch-DDL interface, resource limits, compatibility targets and experimental gate. | + +## Snapshot format acceptance + +Run these cases against actual CLI entry points on local files and S3-compatible +storage. Shared serialization tests alone do not establish lifecycle behavior. + +| Case | Required result | +| --- | --- | +| Version-1 compatibility | Existing Parquet, CSV, JSON and schema-only fixtures retain import, resume and verify behavior, including omitted optional fields. Resuming export preserves version 1. | +| Valid version 2 | Round-trip multiple schemas/chunks, shared packs with different logical schemas, standalone ordinary/large tables, zero-row tables and empty schemas. Check index membership, typed rows and `/`-separated manifest paths on Windows too. | +| Invalid format before DDL | Missing required fields, unsupported manifest/index version or layout, version/layout/format mismatch, non-null checksums, duplicate or unsafe paths, inventory mismatch, invalid ranges, missing/wrong-length objects, and DDL/index membership mismatch fail before any target DDL or data task, including import resume. Ignorable manifest metadata remains accepted. | +| Old reader | Supported release binaries reject a valid version-2 snapshot before target DDL or writes. | +| Export resume | Interrupt object upload, index finalization and manifest completion separately; resume preserves version/layout and completed chunks, cleans only owned incomplete outputs, then produces the same complete inventory. Reject incompatible metadata before cleanup. | +| Import resume | Revalidate metadata before honoring saved state; skip successfully checkpointed tasks, retain DDL completion and replay semantics, and never mark a partially restored task complete. | +| Verify | Valid completed snapshots pass; incomplete chunks, malformed indexes, inventory mismatch, missing/extra files and object-length mismatches fail. A same-length data mutation is not required to fail a structural-only check. | + +## Performance acceptance + +Compare four controls with the same data and concurrency: (A) the existing +per-logical-table export and restore with the explicit-file COPY fix; (B) shared +scans, concurrent writers, redundant export HEAD removal and schema-export +batching, still writing standalone table objects; (C) B plus packing and its +reader, with identical DDL handling; (D) C plus batch DDL. If general restore +request deduplication is available, include it in B as well. + +The proposed release target is at least 3x faster full export and full restore +for D versus A on 10,000 and 100,000 small logical tables with MinIO request +latency increased by 5 ms. Fix this target during RFC review, before running +release measurements. The independent local C-versus-B experiment establishes +neither this comparison nor proof of that target. + +Use Linux release builds and at least three alternating rounds for primary +comparisons. Report medians/spread, schema/data/full durations, CPU, RSS, bytes, +objects, request counts and queue/cache peaks. Include zero-added-latency, +local-file, mixed small/large-table and wide-schema controls. Investigate and +remeasure full-duration regressions exceeding 5% versus A; do not generalize the +latency-injected speedup to these cases. Check values and schemas outside timing. + +Validate actual write/part sizes at or below 8 MiB, aggregate concurrency across +groups/chunks, and resource release after success, cancellation and failure. +Metadata growth and query/codec memory remain visible separately from data +buffer budgets. Report C versus B and D versus C independently; do not multiply +speedups from different workloads or attribute all reader savings to packing. ## Ordering acceptance gate @@ -268,13 +563,13 @@ Exercise both streaming-merge and additional-sort plans. For each case, report frontend and datanode execution plans, peak memory, query memory-accounting metrics, spill volume and elapsed time. Run with the agreed query memory, spill-space, writer and concurrency limits. Verify global -ID order and exactly one complete file per selected logical table in each export -unit, checking schemas and typed rows outside the timed interval. +ID order and exactly one complete indexed Parquet stream per selected logical +table in each schema chunk, checking schemas and typed rows outside timing. Passing requires correct exports within the agreed resource/performance thresholds on both plan paths. Also force resource exhaustion: cancellation must leave the chunk incomplete, and retry with sufficient resources must -produce complete files without altering completed chunks. Writer-only memory -measurements and single-region runs cannot satisfy this gate. Until both the +produce complete indexed data without altering completed chunks. Writer-only +memory measurements and single-region runs cannot satisfy this gate. Until both the normal and failure/retry cases pass, multi-region resource acceptance remains open. diff --git a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md b/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md index 56d75530c7d..9cc7b46464d 100644 --- a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md +++ b/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md @@ -1,109 +1,34 @@ -# Metric export/import experiment results +# Metric export/import experiment conclusions -These experiments support three changes: physical-table export, batched logical -DDL, and removal of repeated directory listing during COPY. They measure separate -parts of export/import; their speedups should not be multiplied. +The experiments support shared physical scans, explicit-file COPY, batched +logical DDL, and packing small logical-table Parquet streams into shared objects. -Measurements were taken on September 10–11, 2026, on an ARM64 machine with 16 GiB -RAM, using local standalone instances, local files, warm caches and unoptimized -test builds. Each table entry represents one run. The results establish useful -optimization directions, not production throughput estimates. +- **Shared scans reduce repeated query work.** Projecting each logical schema + before dictionary expansion preserves its columns and types, but does not + remove the memory cost of a wide physical schema or distributed ordering. +- **Explicit-file COPY removes repeated directory traversal.** This improvement + uses ordinary per-table insertion; it is independent of database merged writes. +- **Packing and its range reader reduce small-object overhead.** They improve + complete export and restore and sharply reduce object-store requests. Benefits + grow with table count and request latency; low-latency export gains are smaller. + Some reader request savings can also benefit standalone files. +- **Batch DDL remains necessary.** Once data restore is faster, logical-table + creation becomes the main remaining restore cost. The batch-DDL experiment used + an internal interface; production transport and retry integration remain work. +- **Bounded writes allow larger objects.** The packing experiment kept writes + and upload parts within 8 MiB without imposing that limit on whole objects or + splitting logical tables. This was the chosen size, not a measured optimum. +- **Packing has a memory trade-off.** CPU and request overhead fell, while peak + process memory increased. Encoding buffers, upload concurrency, range windows + and table/index metadata need separate accounting. -## Physical-table export +The exercised fixtures preserved schemas, DDL and typed values after restore. +Earlier standalone-file experiments used the existing importer; the packing +experiment used an independent driver and prototype range reader. The latter +does not establish production packed-format compatibility or V2 resume safety. -The comparison uses 64 populated logical tables with 1,000 source rows each, -plus an empty table. A time filter selects 63,744 rows, written into 65 logical -Parquet files. Ordinary COPY DATABASE uses parallelism one. Timings cover data -export and exclude schema export and fixture preparation. Both execution orders -were measured at narrow and wide physical schemas. - -| Physical tag columns | Execution order | Ordinary COPY (ms) | Physical-table export (ms) | Speedup | -| ---: | --- | ---: | ---: | ---: | -| 4 | Ordinary COPY first | 889.4 | 262.7 | 3.39x | -| 4 | Physical export first | 893.7 | 309.5 | 2.89x | -| 67 | Ordinary COPY first | 1062.0 | 316.0 | 3.36x | -| 67 | Physical export first | 869.0 | 308.3 | 2.82x | - -Physical-table export was **2.82–3.39x faster** in these runs. Each populated -logical file still contains only its own four columns. Projecting the logical -schema before dictionary expansion kept the expanded-array peak at 0.148 MiB -and the Parquet row-group estimate at 1.112 MiB for both widths. Incoming array -estimates grew from 0.267 MiB to 2.291–2.408 MiB: logical projection does not remove -the memory cost of a wide physical schema. These estimates are not process RSS -limits, and a writer flush threshold does not cap codec allocations. - -The measured single-region query used `SeriesScan` and `SortPreservingMergeExec` -without a separate `SortExec`; `__table_id` matches a physical primary-key prefix. -Only Float64 fields, string tags and timestamps were exercised. These runs do not -establish sorting memory, spill or performance when region count exceeds -frontend `target_partitions`; see the RFC's -[ordering acceptance gate](../2026-09-11-metric-export-import.md#ordering-acceptance-gate). -The timestamp encoder reserved about 1 MiB even for a small batch, illustrating -why a writer flush threshold is not an allocation ceiling. - -## Batched logical DDL - -The restore fixture has logical tables sharing one physical table, one file per -logical table, two rows per populated table and one empty table. The 10,000-table -case contains 19,998 rows. The three paths use SQL HTTP per table, an internal call per table, and -internal batches of up to 128 tables and 1 MiB of rendered SQL, submitted -sequentially. These request bounds do not bound the whole parsed DDL list in -memory. The internal single-call control separates transport savings from -procedure batching. - -| Logical tables | SQL HTTP CREATE (s) | Internal single CREATE (s) | Batched CREATE (s) | Logical calls: single → batch | -| ---: | ---: | ---: | ---: | ---: | -| 100 | 0.318 | 0.177 | 0.056 | 100 → 1 | -| 1,000 | 3.091 | 1.601 | 0.544 | 1,000 → 8 | -| 10,000 | 29.181 | 19.524 | 5.776 | 10,000 → 79 | - -These are logical CREATE call durations. At 10,000 tables, the complete DDL phase -fell from 30.752 s with SQL HTTP to 7.019 s with batching. Total restore fell from -187.279 s to 157.492 s; unchanged data COPY still took about 150 s. Data-phase -variation is not a DDL benefit. Batching improves CREATE itself, but the original -COPY bottleneck limits its effect on total restore. - -The batch executor was in-process. A production authenticated server interface -is still needed; these times do not predict batch HTTP latency. - -## COPY file lookup - -Both variants below use batch DDL and the same COPY concurrency. Previously, -COPY listed a file's parent directory even when it already knew the filename. -With N files in one directory, repeating this for each file causes roughly N² -directory-entry work. The fix checks the known path directly; directory COPY -still lists files normally. - -| Logical tables / files | COPY before (s) | COPY after (s) | Total restore before (s) | Total restore after (s) | -| ---: | ---: | ---: | ---: | ---: | -| 100 | 0.094 | 0.071 | 0.183 | 0.164 | -| 1,000 | 2.135 | 0.612 | 2.849 | 1.289 | -| 10,000 | 146.496 | 5.481 | 153.647 | 12.300 | - -At 10,000 files, COPY improved **26.73x** and total restore improved **12.49x**. -Before the fix, backend/file lookup accounted for about 92% of cumulative active -file-COPY time. A control with profiling disabled completed COPY in 5.449 s. - -The **5.481 s result still uses per-table writes**, with only 19,998 rows. It is -not a merged-write result. Larger files, distributed targets and object stores -are needed to decide whether merged writes provide enough additional benefit. - -## Correctness and remaining coverage - -The experiments compared exported/restored schemas and typed rows with the -source. Checks covered dense/sparse keys, different logical schemas, NULL and -empty values, empty tables, dropped-table residue and fresh target IDs. The -unmodified import-v2 command used in the experiment restored the exported files -into a fresh target with different IDs. This does not establish compatibility -with all released readers. Batch-DDL checks covered two schemas, ordinary -tables/views, failure after a successful 128-table batch and retry of the original -DDL. Existing standalone COPY cases also passed. - -The internal exporter assumed a trusted caller. It could leave a partial output -directory after failure and did not implement completion publication or resume. -These limitations require production integration work. - -Production acceptance still requires release-build measurements with repeated -runs, larger workloads, distributed/multi-region execution, object stores, -released-reader compatibility and export cancellation/recovery. These are tracked -in the [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). +These were local standalone experiments using development builds, local files +and Docker MinIO. Release acceptance still needs the actual packed V2 path, +request-owned reader isolation, old-reader rejection, large-table throughput, +distributed ordering, cancellation and recovery. See the [RFC](../2026-09-11-metric-export-import.md) +and [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). From 16978cf6c2d4aa1f1f6486e6157732bda97e8e76 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 22 Sep 2026 02:39:36 +0000 Subject: [PATCH 45/63] feat(trace): support Semantic Graph for Trace V2 (follow-up to #9192) (#9278) * feat(trace): support Semantic Graph for Trace V2 Signed-off-by: luofucong * fix: remove unused annotation context import Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/common/catalog/src/consts.rs | 6 + src/frontend/src/instance/entity_graph.rs | 105 ++++-- src/frontend/src/instance/jaeger.rs | 7 +- src/operator/src/insert.rs | 2 +- src/operator/src/statement/ddl.rs | 14 +- .../statement/semantic_graph/conventions.rs | 2 +- .../statement/semantic_graph/conventions.yaml | 2 +- src/servers/src/otlp/trace.rs | 7 +- src/table/src/metadata.rs | 1 + src/table/src/requests.rs | 38 +- src/table/src/requests/semantic.rs | 88 ++++- .../system/semantic_graph_trace_v2.result | 329 ++++++++++++++++++ .../common/system/semantic_graph_trace_v2.sql | 165 +++++++++ 13 files changed, 707 insertions(+), 59 deletions(-) create mode 100644 tests/cases/standalone/common/system/semantic_graph_trace_v2.result create mode 100644 tests/cases/standalone/common/system/semantic_graph_trace_v2.sql diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index dd1b5fc3f03..7c71de2e1fa 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -202,6 +202,12 @@ pub fn is_ddl_reserved_table(schema: &str, table: &str) -> bool { pub const TRACE_ID_COLUMN: &str = "trace_id"; pub const SPAN_ID_COLUMN: &str = "span_id"; pub const SPAN_NAME_COLUMN: &str = "span_name"; +/// Column containing span attributes as JSON/JSON2. +pub const SPAN_ATTRIBUTES_COLUMN: &str = "span_attributes"; +/// Column containing instrumentation scope attributes as JSON2. +pub const SCOPE_ATTRIBUTES_COLUMN: &str = "scope_attributes"; +/// Column containing resource attributes as JSON2. +pub const RESOURCE_ATTRIBUTES_COLUMN: &str = "resource_attributes"; pub const SERVICE_NAME_COLUMN: &str = "service_name"; pub const PARENT_SPAN_ID_COLUMN: &str = "parent_span_id"; // More fixed columns/values of the `greptime_trace_v1` data model, shared by diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 6310d0614dd..09c23651bcc 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -23,8 +23,8 @@ //! the query engine. Injected into the catalog manager after the engine is built, //! breaking the `catalog -> query` cycle. -use std::collections::{BTreeMap, HashMap}; -use std::sync::Weak; +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::{Arc, Weak}; use async_trait::async_trait; use auth::{ @@ -41,13 +41,16 @@ use common_catalog::consts::{ }; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; +use common_function::scalars::json::json_get::JsonGetWithType; +use common_function::scalars::udf::create_udf; use common_query::OutputData; use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; use common_recordbatch::SendableRecordBatchStream; use common_telemetry::{debug, warn}; use common_time::timestamp::TimeUnit; +use datafusion::common::{Column, ScalarValue}; use datafusion::dataframe::DataFrame; -use datafusion_expr::LogicalPlan; +use datafusion_expr::{Expr, LogicalPlan, lit}; use futures::TryStreamExt; use operator::statement::semantic_graph::{ CallsSource, CoDeclaredSource, Conventions, DeclaredSource, ENTITY_TYPE_GEN_AI_AGENT, @@ -64,8 +67,8 @@ use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, - SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns, - parse_entity_option_key, + SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, TABLE_DATA_MODEL_TRACE_V2, is_trace_table, + parse_entity_columns, parse_entity_option_key, trace_v2_attribute, }; use crate::error; @@ -164,12 +167,11 @@ impl EntityGraphProviderImpl { |(entity_type, (id_columns, descriptive_columns, scope_columns))| { // A stale declaration (e.g. its column was dropped later) // must not poison every graph scan; skip it. - let schema = &table_info.meta.schema; if let Some(missing) = id_columns .iter() .chain(&descriptive_columns) .chain(&scope_columns) - .find(|c| schema.column_schema_by_name(c).is_none()) + .find(|c| !Self::can_resolve_entity_column(table_info, c)) { warn!( "Skipping entity declaration `{}` of table `{}`: column `{}` not found", @@ -194,7 +196,7 @@ impl EntityGraphProviderImpl { } /// All entity declarations of one table: the explicit options plus the - /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 + /// zero-configuration conventions (`otlp_trace_entities` for trace-v1/v2 /// tables — including the `service` identity of tables created before the /// ingest-side auto-stamp — and the Prometheus/OTel descriptor /// whitelists). An explicit declaration of a type always suppresses the @@ -207,7 +209,7 @@ impl EntityGraphProviderImpl { ) -> Vec { let mut declarations = Self::parse_declarations(table_info); let mut supersessions = Vec::new(); - if is_trace_v1_table(table_info) { + if is_trace_table(table_info) { Self::extend_with_implicit_entities( table_info, &conventions.otlp_trace_entities, @@ -303,8 +305,8 @@ impl EntityGraphProviderImpl { } /// Synthesizes the applicable subset of `entities` on `table_info`: - /// explicit declarations win, every id column must exist (no guessing), - /// descriptive columns are filtered to those present. + /// explicit declarations win, every id reference must be resolvable, + /// and descriptive references are filtered by the same rule. fn extend_with_implicit_entities( table_info: &TableInfo, entities: &[ImplicitEntity], @@ -330,7 +332,7 @@ impl EntityGraphProviderImpl { if let Some(missing) = implicit .id .iter() - .find(|c| schema.column_schema_by_name(c).is_none()) + .find(|c| !Self::can_resolve_entity_column(table_info, c)) { debug!( "Table `{}` lacks the id column `{}`; skipping the implicit `{}` declaration", @@ -350,7 +352,7 @@ impl EntityGraphProviderImpl { implicit .descriptive .iter() - .filter(|c| schema.column_schema_by_name(c).is_some()) + .filter(|c| Self::can_resolve_entity_column(table_info, c)) .cloned() .collect() }; @@ -359,7 +361,7 @@ impl EntityGraphProviderImpl { let id_qualifier = implicit .qualified_by .clone() - .filter(|c| schema.column_schema_by_name(c).is_some()); + .filter(|c| Self::can_resolve_entity_column(table_info, c)); if let Some(entity_type) = &implicit.superseded_by { supersessions.push((declarations.len(), entity_type.clone())); } @@ -434,7 +436,7 @@ impl EntityGraphProviderImpl { while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { let table_info = table.table_info(); let table_declarations = Self::declarations_for(&table_info, conventions); - let is_trace = is_trace_v1_table(&table_info); + let is_trace = is_trace_table(&table_info); // Authorize only tables that would contribute rows. if per_table_auth && (is_trace || !table_declarations.is_empty()) @@ -532,8 +534,63 @@ impl EntityGraphProviderImpl { } } - fn read_table(&self, table: TableRef) -> Result { - self.query_engine.read_table(table).map_err(BoxedError::new) + fn can_resolve_entity_column(table_info: &TableInfo, column: &str) -> bool { + let table_schema = &table_info.meta.schema; + if table_schema.column_schema_by_name(column).is_some() { + return true; + } + let data_model = table_info.meta.options.data_model(); + trace_v2_attribute(table_schema, data_model, column).is_some() + } + + /// Projects referenced V2 attributes under their declaration names for + /// the shared entity and relationship plans. Other models scan unchanged. + fn read_table( + &self, + table: TableRef, + declarations: &[EntityDeclaration], + ) -> Result { + let info = table.table_info(); + let mut scan = self + .query_engine + .read_table(table) + .map_err(BoxedError::new)?; + let data_model = info.meta.options.data_model(); + if data_model != Some(TABLE_DATA_MODEL_TRACE_V2) { + return Ok(scan); + } + + // Read Trace V2 table: + let mut columns = BTreeSet::new(); + for declaration in declarations { + columns.extend(declaration.id_columns.iter()); + columns.extend(declaration.descriptive_columns.iter()); + columns.extend(declaration.scope_columns.iter()); + columns.extend(declaration.id_qualifier.iter()); + columns.extend(declaration.superseded_by_columns.iter()); + } + let conventions = conventions() + .map_err(datafusion::error::DataFusionError::Internal) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)?; + let virtual_candidates = &conventions.virtual_dst_candidates; + columns.extend(virtual_candidates.iter().map(|candidate| &candidate.column)); + let get = create_udf(Arc::new(JsonGetWithType::default())); + for column in columns { + if let Some((root, key)) = trace_v2_attribute(&info.meta.schema, data_model, column) { + let path = format!("$.{}", serde_json::Value::String(key.to_string())); + let value = get.call(vec![ + Expr::Column(Column::from_name(root)), + lit(path), + lit(ScalarValue::Utf8View(None)), + ]); + scan = scan + .with_column(column, value) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)?; + } + } + Ok(scan) } /// The declared-edge branch source, when the physical table exists, the @@ -587,7 +644,7 @@ impl EntityGraphProviderImpl { )); } Ok(Some(DeclaredSource { - scan: self.read_table(table)?, + scan: self.read_table(table, &[])?, })) } @@ -633,8 +690,8 @@ impl EntityGraphProvider for EntityGraphProviderImpl { let mut plans = Vec::with_capacity(sources.len()); for source in sources { plans.push(RegistrySource { + scan: self.read_table(source.table, &source.declarations)?, declarations: source.declarations, - scan: self.read_table(source.table)?, }); } let Some(window) = Self::query_window(&request)? else { @@ -658,18 +715,24 @@ impl EntityGraphProvider for EntityGraphProviderImpl { let (sources, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; let mut calls = Vec::with_capacity(traces.len()); for trace in traces { + let declarations = trace + .service + .iter() + .chain(trace.agent.iter()) + .cloned() + .collect::>(); calls.push(CallsSource { + scan: self.read_table(trace.table, &declarations)?, service: trace.service, agent: trace.agent, - scan: self.read_table(trace.table)?, }); } let mut co_declared = Vec::with_capacity(sources.len()); for source in sources { co_declared.push(CoDeclaredSource { + scan: self.read_table(source.table, &source.declarations)?, declarations: source.declarations, is_trace: source.is_trace, - scan: self.read_table(source.table)?, }); } let declared = self.declared_source(catalog, query_ctx.as_deref()).await?; diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index 556f7640de8..c9b5bd9e988 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -370,12 +370,7 @@ async fn query_trace_table( })?; let table_info = table.table_info(); - let data_model = table_info - .meta - .options - .extra_options - .get(table::requests::TABLE_DATA_MODEL) - .map(String::as_str); + let data_model = table_info.meta.options.data_model(); // collect to set let col_names = table_info diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 989a4aa399d..afc903621e2 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -1509,7 +1509,7 @@ pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) else { return Ok(()); }; - if let Some(actual) = table_info.meta.options.extra_options.get(TABLE_DATA_MODEL) { + if let Some(actual) = table_info.meta.options.data_model() { ensure!( actual == expected, InvalidInsertRequestSnafu { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 0bebd3c1295..772bff4aad9 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -100,8 +100,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, AnnotationContext, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, - INGEST_ROWS_RATE_LIMIT_KEY, TableOptions, validate_and_normalize_annotation_options, + AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, INGEST_ROWS_RATE_LIMIT_KEY, + TableOptions, validate_and_normalize_annotation_options, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -2657,11 +2657,7 @@ fn validate_and_normalize_annotations( partition_key_indices: &[usize], ) -> Result<()> { use table::requests::AnnotationValidationError as CheckError; - let cx = AnnotationContext { - schema, - partition_key_indices, - }; - validate_and_normalize_annotation_options(options, &cx).map_err(|e| match e { + let handle_error = |e: CheckError| match e { CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(), e @ (CheckError::UnknownKey { .. } | CheckError::InvalidValue { .. } @@ -2676,7 +2672,9 @@ fn validate_and_normalize_annotations( reason: e.to_string(), } .build(), - }) + }; + validate_and_normalize_annotation_options(options, schema, partition_key_indices) + .map_err(handle_error) } fn find_partition_columns(partitions: &Option) -> Result> { diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs index 572140c2064..7120f3102f1 100644 --- a/src/operator/src/statement/semantic_graph/conventions.rs +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -44,7 +44,7 @@ pub struct VirtualDstCandidate { } /// One implicit entity declaration: of a whitelisted Prometheus or OTel info -/// metric, or of a trace-v1 table's flattened resource attributes. +/// metric, or of a Trace V1/V2 table's resource attributes. #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct ImplicitEntity { diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml index 85dfe932aba..4d9f05ed519 100644 --- a/src/operator/src/statement/semantic_graph/conventions.yaml +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -37,7 +37,7 @@ virtual_dst_candidates: - { column: span_attributes.db.name, connection_type: database } - { column: span_attributes.server.address, connection_type: virtual_node } -# Implicit declarations for greptime_trace_v1 tables; the pod UID bridges +# Implicit declarations for Trace V1/V2 tables; the pod UID bridges # trace-side pods onto the kube-state-metrics entities below. host.name and # container.name vary by SDK and resource detector, so identity is the id. # qualified_by folds the namespace into the service identity the same way the diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index 60865472086..73ce493388b 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -28,7 +28,8 @@ pub use common_catalog::consts::{ SPAN_STATUS_CODE_COLUMN as SPAN_STATUS_CODE, TRACE_TIMESTAMP_COLUMN as TIMESTAMP_COLUMN, }; pub use common_catalog::consts::{ - PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_NAME_COLUMN, TRACE_ID_COLUMN, + PARENT_SPAN_ID_COLUMN, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, + SPAN_ATTRIBUTES_COLUMN, SPAN_ID_COLUMN, SPAN_NAME_COLUMN, TRACE_ID_COLUMN, }; use pipeline::PipelineWay; @@ -37,14 +38,10 @@ use crate::otlp::trace::span::TraceSpan; pub const TIMESTAMP_END_COLUMN: &str = "timestamp_end"; pub const SPAN_STATUS_MESSAGE_COLUMN: &str = "span_status_message"; -pub const SPAN_ATTRIBUTES_COLUMN: &str = "span_attributes"; pub const SPAN_EVENTS_COLUMN: &str = "span_events"; pub const SPAN_LINKS_COLUMN: &str = "span_links"; pub const SCOPE_NAME_COLUMN: &str = "scope_name"; pub const SCOPE_VERSION_COLUMN: &str = "scope_version"; -/// Column containing instrumentation scope attributes as a JSON2 object. -pub const SCOPE_ATTRIBUTES_COLUMN: &str = "scope_attributes"; -pub const RESOURCE_ATTRIBUTES_COLUMN: &str = "resource_attributes"; pub const TRACE_STATE_COLUMN: &str = "trace_state"; // const keys diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index a784fbcf86b..5916dd7ab49 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -470,6 +470,7 @@ impl TableMeta { .build() })?; let cx = AnnotationContext { + data_model: self.options.data_model(), schema: &self.schema, partition_key_indices: &self.partition_key_indices, }; diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 10055df4b45..1ca99268d67 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -63,17 +63,14 @@ pub const TABLE_DATA_MODEL_TRACE_V1: &str = "greptime_trace_v1"; /// Table data model used by the JSON2-based OTLP trace pipeline. pub const TABLE_DATA_MODEL_TRACE_V2: &str = "greptime_trace_v2"; -/// Returns true if the table stores spans in the `greptime_trace_v1` data model -/// (fixed span columns), the shape the Jaeger query path and the entity-graph -/// derivation rely on. -pub fn is_trace_v1_table(table_info: &crate::metadata::TableInfo) -> bool { - table_info - .meta - .options - .extra_options - .get(TABLE_DATA_MODEL) - .map(|v| v == TABLE_DATA_MODEL_TRACE_V1) - .unwrap_or(false) +/// Returns true for the Trace V1 and V2 data models supported +/// by semantic graph derivation. +pub fn is_trace_table(table_info: &crate::metadata::TableInfo) -> bool { + let table_data_model = table_info.meta.options.data_model(); + matches!( + table_data_model, + Some(TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2) + ) } pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat"; @@ -226,6 +223,11 @@ pub const REPARTITION_COLUMN_HINT_KEY: &str = "repartition.column.hint"; pub const REPARTITION_PARTITION_NUM_HINT_KEY: &str = "repartition.partition.num.hint"; impl TableOptions { + /// Returns the table data model, if specified. + pub fn data_model(&self) -> Option<&str> { + self.extra_options.get(TABLE_DATA_MODEL).map(String::as_str) + } + pub fn try_from_iter>( iter: U, ) -> Result { @@ -480,6 +482,7 @@ pub fn validate_annotation_keys<'a>( /// Table shape an annotation option is validated against. pub struct AnnotationContext<'a> { + pub data_model: Option<&'a str>, pub schema: &'a Schema, pub partition_key_indices: &'a [usize], } @@ -570,6 +573,9 @@ pub(crate) fn validate_and_normalize_annotation( } if parse_entity_option_key(key).is_some() { for column in parse_entity_columns(value) { + if trace_v2_attribute(cx.schema, cx.data_model, &column).is_some() { + continue; + } let schema = cx.schema.column_schema_by_name(&column).ok_or_else(|| { AnnotationValidationError::ColumnNotFound { column: column.clone(), @@ -620,14 +626,20 @@ pub(crate) fn validate_and_normalize_annotation( /// and writes normalized values back in place. pub fn validate_and_normalize_annotation_options( options: &mut TableOptions, - cx: &AnnotationContext<'_>, + schema: &Schema, + partition_key_indices: &[usize], ) -> std::result::Result<(), AnnotationValidationError> { + let cx = AnnotationContext { + data_model: options.data_model(), + schema, + partition_key_indices, + }; let mut normalized = Vec::new(); for (key, value) in &options.extra_options { let Some(family) = AnnotationFamily::of_key(key) else { continue; }; - let checked = validate_and_normalize_annotation(family, cx, key, value)?; + let checked = validate_and_normalize_annotation(family, &cx, key, value)?; if checked != *value { normalized.push((key.clone(), checked)); } diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 7c51a1551fd..dbe28eb240f 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -30,7 +30,13 @@ //! [`crate::requests::validate_table_option`], so they are accepted both on the //! ingestion auto-create path and on explicit `CREATE TABLE ... WITH (...)` DDL. +use common_catalog::consts::{ + RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, SPAN_ATTRIBUTES_COLUMN, +}; use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; + +use crate::requests::TABLE_DATA_MODEL_TRACE_V2; /// Reserved prefix for every public semantic table-option key. pub const SEMANTIC_PREFIX: &str = "greptime.semantic."; @@ -193,6 +199,34 @@ pub fn is_entity_option_key(key: &str) -> bool { parse_entity_option_key(key).is_some() } +/// Resolves a Trace V2 entity reference to a JSON2 root and a literal attribute key. +/// Everything after the first dot belongs to the key, matching Trace V1's +/// flattened column names. Existing physical columns take precedence. +pub fn trace_v2_attribute<'a>( + schema: &Schema, + data_model: Option<&str>, + column: &'a str, +) -> Option<(&'a str, &'a str)> { + if data_model != Some(TABLE_DATA_MODEL_TRACE_V2) + || schema.column_schema_by_name(column).is_some() + { + return None; + } + let (root, key) = column.split_once('.')?; + (!key.is_empty() + && schema + .column_schema_by_name(root) + .is_some_and(is_v2_attribute_column)) + .then_some((root, key)) +} + +fn is_v2_attribute_column(schema: &ColumnSchema) -> bool { + matches!( + schema.name.as_str(), + RESOURCE_ATTRIBUTES_COLUMN | SCOPE_ATTRIBUTES_COLUMN | SPAN_ATTRIBUTES_COLUMN + ) && schema.data_type.is_json2() +} + /// Returns true if a column of `data_type` renders as a stable string — the /// requirement for entity id/descriptive/scope columns. The read-time /// derivation casts them to strings, so a type without a stable string form @@ -217,8 +251,9 @@ pub fn has_stable_string_form(data_type: &ConcreteDataType) -> bool { pub fn parse_entity_columns(value: &str) -> Vec { value .split(',') - .map(|c| c.trim().to_string()) - .filter(|c| !c.is_empty()) + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(String::from) .collect() } @@ -244,7 +279,7 @@ pub fn is_semantic_option_key(key: &str) -> bool { /// are rejected. pub fn validate_semantic_option(key: &str, value: &str) -> bool { if is_entity_option_key(key) { - return !value.is_empty() && value.split(',').all(|col| !col.trim().is_empty()); + return !value.is_empty() && value.split(',').all(|column| !column.trim().is_empty()); } match key { SEMANTIC_PIPELINE @@ -293,8 +328,55 @@ pub fn validate_semantic_option(key: &str, value: &str) -> bool { #[cfg(test)] mod tests { + use datatypes::types::JsonType; + use super::*; + #[test] + fn test_trace_v2_attribute_references() { + let schema = Schema::new(vec![ + ColumnSchema::new( + RESOURCE_ATTRIBUTES_COLUMN, + ConcreteDataType::Json(JsonType::null()), + true, + ), + ColumnSchema::new( + SCOPE_ATTRIBUTES_COLUMN, + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "resource_attributes.physical", + ConcreteDataType::string_datatype(), + true, + ), + ]); + let model = Some(TABLE_DATA_MODEL_TRACE_V2); + for key in ["host.name", "host..name", "host[0]", "host.*", r#"a"b\c.d"#] { + let column = format!("resource_attributes.{key}"); + assert_eq!( + trace_v2_attribute(&schema, model, &column), + Some((RESOURCE_ATTRIBUTES_COLUMN, key)) + ); + } + for column in [ + "resource_attributes", + "resource_attributes.", + "scope_attributes.name", + "span_attributes.name", + "resource_attributes.physical", + "missing.key", + ] { + assert_eq!(trace_v2_attribute(&schema, model, column), None, "{column}"); + } + for model in [None, Some(crate::requests::TABLE_DATA_MODEL_TRACE_V1)] { + assert_eq!( + trace_v2_attribute(&schema, model, "resource_attributes.host.name"), + None + ); + } + } + #[test] fn test_is_semantic_option_key() { assert!(is_semantic_option_key(SEMANTIC_SIGNAL_TYPE)); diff --git a/tests/cases/standalone/common/system/semantic_graph_trace_v2.result b/tests/cases/standalone/common/system/semantic_graph_trace_v2.result new file mode 100644 index 00000000000..753f273d59a --- /dev/null +++ b/tests/cases/standalone/common/system/semantic_graph_trace_v2.result @@ -0,0 +1,329 @@ +-- Trace V2 uses the same flattened attribute references as V1. +-- Everything after the root column is one literal key, including dots. +create table graph_trace_v2 ( + "timestamp" timestamp(9) time index, + timestamp_end timestamp(9), + duration_nano bigint, + parent_span_id string, + trace_id string, + span_id string, + span_kind string, + span_name string, + span_status_code string, + span_status_message string, + trace_state string, + scope_name string, + scope_version string, + service_name string, + span_attributes json2, + scope_attributes json2, + resource_attributes json2, + span_events json, + span_links json, + primary key (service_name) +) with ( + 'table_data_model' = 'greptime_trace_v2', + 'append_mode' = 'true', + 'greptime.semantic.entity.gen_ai.agent.id' = 'span_attributes.gen_ai.agent.id', + 'greptime.semantic.entity.gen_ai.agent.scope' = 'scope_attributes.environment', + 'greptime.semantic.entity.gen_ai.model.id' = 'span_attributes.gen_ai.request.model', + 'greptime.semantic.entity.literal_peer.id' = 'span_attributes.db.namespace', + 'greptime.semantic.entity.literal_resource.id' = 'resource_attributes.host.name', + 'greptime.semantic.entity.literal_span.id' = 'span_attributes.gen_ai.agent.id', + 'greptime.semantic.entity.literal_scope.id' = 'scope_attributes.deployment.region.name' +); + +Affected Rows: 0 + +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, span_kind, span_status_code, + service_name, span_attributes, scope_attributes, resource_attributes) +values + ('2026-01-01 00:00:01', 100000000, 'paired', 'client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', + 'frontend', '{"gen_ai.agent.id":"agent-a","gen_ai.request.model":"model-a","gen_ai":{"agent":{"id":"nested-agent"}}}', + '{"environment":"prod","deployment.region.name":"literal-region","deployment":{"region.name":"nested-region"},"a\"b\\c.d":7}', '{"service.namespace":"shop","service.instance.id":"front-1","host.id":"host-1","host.name":"node-1","host":{"name":"nested-node"},"service":{"namespace":"wrong-namespace"}}'), + ('2026-01-01 00:00:02', 200000000, 'virtual', 'client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_ERROR', + 'frontend', '{"service.peer.name":"","db.namespace":"orders","db":{"namespace":"wrong-orders"}}', '{}', + '{"service.namespace":"shop","service.instance.id":"front-1","host.id":"host-1","host.name":"node-1"}'), + ('2026-01-01 00:00:03', 100, 'missing', 'internal', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', + 'no-resource', '{"gen_ai.agent.id":null}', '{}', '{}'); + +Affected Rows: 3 + +-- Agent parent/child references are independent of span kind. +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, parent_span_id, span_kind, + span_status_code, service_name, span_attributes, scope_attributes, resource_attributes) +values ('2026-01-01 00:00:01', 30000000, 'paired', 'agent-child', 'client', + 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'frontend', + '{"gen_ai.agent.id":"agent-b"}', '{"environment":"prod"}', '{"service.namespace":"shop"}'); + +Affected Rows: 1 + +-- Calls can pair a V2 client with a V1 server. The literal_resource declaration +-- has exactly the same spelling in both models, including its output JSON keys. +create table graph_trace_v1_peer ( + "timestamp" timestamp(9) time index, + trace_id string, span_id string, parent_span_id string, + span_kind string, span_status_code string, duration_nano bigint, + service_name string, + "resource_attributes.service.namespace" string, + "resource_attributes.host.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', + 'greptime.semantic.entity.literal_resource.id' = 'resource_attributes.host.name'); + +Affected Rows: 0 + +insert into graph_trace_v1_peer values + ('2026-01-01 00:00:01', 'paired', 'server', 'client', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 50000000, 'backend', 'shop', 'node-v1'); + +Affected Rows: 1 + +select entity_type, entity_id, scope +from greptime_private.semantic_entities +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + ++------------------+-----------------------+-------+ +| entity_type | entity_id | scope | ++------------------+-----------------------+-------+ +| gen_ai.agent | agent-a | prod | +| gen_ai.agent | agent-b | prod | +| gen_ai.model | model-a | | +| host | host-1 | | +| literal_peer | orders | | +| literal_resource | node-1 | | +| literal_resource | node-v1 | | +| literal_scope | literal-region | | +| literal_span | agent-a | | +| literal_span | agent-b | | +| service | no-resource | | +| service | shop/backend | | +| service | shop/frontend | | +| service.instance | shop/frontend,front-1 | | ++------------------+-----------------------+-------+ + +select src_type, src_id, dst_type, dst_id, rel_type, request_count, error_count +from greptime_private.semantic_relationships +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, src_type, src_id, dst_id; + ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ +| src_type | src_id | dst_type | dst_id | rel_type | request_count | error_count | ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ +| gen_ai.agent | agent-a | gen_ai.agent | agent-b | calls | 1 | 0 | +| service | shop/frontend | service | orders | calls | 1 | 1 | +| service | shop/frontend | service | shop/backend | calls | 1 | 0 | +| service.instance | shop/frontend,front-1 | service | shop/frontend | part_of | | | +| service.instance | shop/frontend,front-1 | host | host-1 | runs_on | | | +| gen_ai.agent | agent-a | gen_ai.model | model-a | uses | | | ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ + +-- Empty components in the V1-compatible comma-separated list remain invalid. +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.host.id,'; + +Error: 1004(InvalidArguments), Invalid alter table(graph_trace_v2) request: invalid value `resource_attributes.host.id,` for semantic option `greptime.semantic.entity.invalid.id` + +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.host.id,,service_name'; + +Error: 1004(InvalidArguments), Invalid alter table(graph_trace_v2) request: invalid value `resource_attributes.host.id,,service_name` for semantic option `greptime.semantic.entity.invalid.id` + +-- ALTER accepts JSON2 references; a missing key simply yields no entity. +alter table graph_trace_v2 set 'greptime.semantic.entity.custom.id' = 'scope_attributes.missing'; + +Affected Rows: 0 + +alter table graph_trace_v2 set 'greptime.semantic.entity.host.descriptive' = 'resource_attributes.host.name'; + +Affected Rows: 0 + +-- Preserve implicit host identity by explicitly declaring it as well. +alter table graph_trace_v2 set 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'; + +Affected Rows: 0 + +admin flush_table('graph_trace_v2'); + ++-------------------------------------+ +| ADMIN flush_table('graph_trace_v2') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +admin flush_table('graph_trace_v1_peer'); + ++------------------------------------------+ +| ADMIN flush_table('graph_trace_v1_peer') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, scope, entity_id_attrs, descriptive +from greptime_private.semantic_entities +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + ++------------------+-----------------------+-------+--------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------+ +| entity_type | entity_id | scope | entity_id_attrs | descriptive | ++------------------+-----------------------+-------+--------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------+ +| gen_ai.agent | agent-a | prod | {"span_attributes.gen_ai.agent.id":"agent-a"} | | +| gen_ai.agent | agent-b | prod | {"span_attributes.gen_ai.agent.id":"agent-b"} | | +| gen_ai.model | model-a | | {"span_attributes.gen_ai.request.model":"model-a"} | | +| host | host-1 | | {"resource_attributes.host.id":"host-1"} | {"resource_attributes.host.name":"node-1"} | +| literal_peer | orders | | {"span_attributes.db.namespace":"orders"} | | +| literal_resource | node-1 | | {"resource_attributes.host.name":"node-1"} | | +| literal_resource | node-v1 | | {"resource_attributes.host.name":"node-v1"} | | +| literal_scope | literal-region | | {"scope_attributes.deployment.region.name":"literal-region"} | | +| literal_span | agent-a | | {"span_attributes.gen_ai.agent.id":"agent-a"} | | +| literal_span | agent-b | | {"span_attributes.gen_ai.agent.id":"agent-b"} | | +| service | no-resource | | {"resource_attributes.service.namespace":"","service_name":"no-resource"} | | +| service | shop/backend | | {"resource_attributes.service.namespace":"shop","service_name":"backend"} | | +| service | shop/frontend | | {"resource_attributes.service.namespace":"shop","service_name":"frontend"} | | +| service.instance | shop/frontend,front-1 | | {"resource_attributes.service.instance.id":"front-1","resource_attributes.service.namespace":"shop","service_name":"frontend"} | | ++------------------+-----------------------+-------+--------------------------------------------------------------------------------------------------------------------------------+--------------------------------------------+ + +select src_type, src_id, dst_type, dst_id, rel_type, request_count, error_count +from greptime_private.semantic_relationships +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, src_type, src_id, dst_id; + ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ +| src_type | src_id | dst_type | dst_id | rel_type | request_count | error_count | ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ +| gen_ai.agent | agent-a | gen_ai.agent | agent-b | calls | 1 | 0 | +| service | shop/frontend | service | orders | calls | 1 | 1 | +| service | shop/frontend | service | shop/backend | calls | 1 | 0 | +| service.instance | shop/frontend,front-1 | service | shop/frontend | part_of | | | +| service.instance | shop/frontend,front-1 | host | host-1 | runs_on | | | +| gen_ai.agent | agent-a | gen_ai.model | model-a | uses | | | ++------------------+-----------------------+--------------+---------------+----------+---------------+-------------+ + +-- Quotes, backslashes and dots remain a literal key; numbers render as strings. +alter table graph_trace_v2 set 'greptime.semantic.entity.custom.id' = 'scope_attributes.a"b\c.d'; + +Affected Rows: 0 + +select entity_type, entity_id +from greptime_private.semantic_entities +where entity_type = 'custom' + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00'; + ++-------------+-----------+ +| entity_type | entity_id | ++-------------+-----------+ +| custom | 7 | ++-------------+-----------+ + +-- Path-looking suffixes are literal keys, not nested traversal or array access. +-- These missing keys are valid declarations and produce no entities. +alter table graph_trace_v2 set 'greptime.semantic.entity.literal_missing.id' = 'resource_attributes.host..name'; + +Affected Rows: 0 + +alter table graph_trace_v2 set 'greptime.semantic.entity.literal_index.id' = 'resource_attributes.host[0]'; + +Affected Rows: 0 + +-- The root must exist and be one of the three JSON2 attribute columns. +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'span_events.name'; + +Error: 4002(TableColumnNotFound), Column span_events.name not exists in table graph_trace_v2 + +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.'; + +Error: 4002(TableColumnNotFound), Column resource_attributes. not exists in table graph_trace_v2 + +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'missing.key'; + +Error: 4002(TableColumnNotFound), Column missing.key not exists in table graph_trace_v2 + +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes'; + +Error: 1004(InvalidArguments), Invalid alter table(graph_trace_v2) request: entity column `resource_attributes` (option `greptime.semantic.entity.invalid.id`) has type `Json2""`, which cannot render as a string + +create table graph_bad_v2 ( + ts timestamp time index, resource_attributes string +) with ('table_data_model' = 'greptime_trace_v2', + 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'); + +Error: 1004(InvalidArguments), Cannot find column by name: resource_attributes.host.id + +-- Ordinary JSON2 tables do not acquire Trace V2 attribute-reference semantics. +create table graph_not_trace ( + ts timestamp time index, resource_attributes json2 +) with ('append_mode' = 'true', + 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'); + +Error: 1004(InvalidArguments), Cannot find column by name: resource_attributes.host.id + +-- Physical convention columns win over conflicting JSON2 keys. +alter table graph_trace_v2 add column "resource_attributes.host.id" string; + +Affected Rows: 0 + +alter table graph_trace_v2 add column "resource_attributes.service.namespace" string; + +Affected Rows: 0 + +alter table graph_trace_v2 add column "span_attributes.db.namespace" string; + +Affected Rows: 0 + +-- Restore the implicit host declaration for the physical-column check. +alter table graph_trace_v2 unset 'greptime.semantic.entity.host.id'; + +Affected Rows: 0 + +alter table graph_trace_v2 unset 'greptime.semantic.entity.host.descriptive'; + +Affected Rows: 0 + +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, span_kind, span_status_code, + service_name, resource_attributes, span_attributes, + "resource_attributes.host.id", "resource_attributes.service.namespace", "span_attributes.db.namespace") +values ('2026-01-01 00:00:04', 100, 'physical', 'physical-client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', + 'physical-service', '{"host.id":"json-host","service.namespace":"json-namespace"}', + '{"db.namespace":"json-database","db":{"namespace":"nested-database"}}', + 'physical-host', 'physical-namespace', 'physical-database'); + +Affected Rows: 1 + +select entity_type, entity_id +from greptime_private.semantic_entities +where entity_type in ('host', 'service') + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + ++-------------+-------------------------------------+ +| entity_type | entity_id | ++-------------+-------------------------------------+ +| host | physical-host | +| service | frontend | +| service | no-resource | +| service | physical-namespace/physical-service | +| service | shop/backend | ++-------------+-------------------------------------+ + +select src_id, dst_id, rel_type +from greptime_private.semantic_relationships +where src_id = 'physical-namespace/physical-service' + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, dst_id; + ++-------------------------------------+-------------------+----------+ +| src_id | dst_id | rel_type | ++-------------------------------------+-------------------+----------+ +| physical-namespace/physical-service | physical-database | calls | ++-------------------------------------+-------------------+----------+ + +drop table graph_trace_v2; + +Affected Rows: 0 + +drop table graph_trace_v1_peer; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph_trace_v2.sql b/tests/cases/standalone/common/system/semantic_graph_trace_v2.sql new file mode 100644 index 00000000000..3c9ca25150f --- /dev/null +++ b/tests/cases/standalone/common/system/semantic_graph_trace_v2.sql @@ -0,0 +1,165 @@ +-- Trace V2 uses the same flattened attribute references as V1. +-- Everything after the root column is one literal key, including dots. +create table graph_trace_v2 ( + "timestamp" timestamp(9) time index, + timestamp_end timestamp(9), + duration_nano bigint, + parent_span_id string, + trace_id string, + span_id string, + span_kind string, + span_name string, + span_status_code string, + span_status_message string, + trace_state string, + scope_name string, + scope_version string, + service_name string, + span_attributes json2, + scope_attributes json2, + resource_attributes json2, + span_events json, + span_links json, + primary key (service_name) +) with ( + 'table_data_model' = 'greptime_trace_v2', + 'append_mode' = 'true', + 'greptime.semantic.entity.gen_ai.agent.id' = 'span_attributes.gen_ai.agent.id', + 'greptime.semantic.entity.gen_ai.agent.scope' = 'scope_attributes.environment', + 'greptime.semantic.entity.gen_ai.model.id' = 'span_attributes.gen_ai.request.model', + 'greptime.semantic.entity.literal_peer.id' = 'span_attributes.db.namespace', + 'greptime.semantic.entity.literal_resource.id' = 'resource_attributes.host.name', + 'greptime.semantic.entity.literal_span.id' = 'span_attributes.gen_ai.agent.id', + 'greptime.semantic.entity.literal_scope.id' = 'scope_attributes.deployment.region.name' +); + +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, span_kind, span_status_code, + service_name, span_attributes, scope_attributes, resource_attributes) +values + ('2026-01-01 00:00:01', 100000000, 'paired', 'client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', + 'frontend', '{"gen_ai.agent.id":"agent-a","gen_ai.request.model":"model-a","gen_ai":{"agent":{"id":"nested-agent"}}}', + '{"environment":"prod","deployment.region.name":"literal-region","deployment":{"region.name":"nested-region"},"a\"b\\c.d":7}', '{"service.namespace":"shop","service.instance.id":"front-1","host.id":"host-1","host.name":"node-1","host":{"name":"nested-node"},"service":{"namespace":"wrong-namespace"}}'), + ('2026-01-01 00:00:02', 200000000, 'virtual', 'client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_ERROR', + 'frontend', '{"service.peer.name":"","db.namespace":"orders","db":{"namespace":"wrong-orders"}}', '{}', + '{"service.namespace":"shop","service.instance.id":"front-1","host.id":"host-1","host.name":"node-1"}'), + ('2026-01-01 00:00:03', 100, 'missing', 'internal', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', + 'no-resource', '{"gen_ai.agent.id":null}', '{}', '{}'); + +-- Agent parent/child references are independent of span kind. +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, parent_span_id, span_kind, + span_status_code, service_name, span_attributes, scope_attributes, resource_attributes) +values ('2026-01-01 00:00:01', 30000000, 'paired', 'agent-child', 'client', + 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'frontend', + '{"gen_ai.agent.id":"agent-b"}', '{"environment":"prod"}', '{"service.namespace":"shop"}'); + +-- Calls can pair a V2 client with a V1 server. The literal_resource declaration +-- has exactly the same spelling in both models, including its output JSON keys. +create table graph_trace_v1_peer ( + "timestamp" timestamp(9) time index, + trace_id string, span_id string, parent_span_id string, + span_kind string, span_status_code string, duration_nano bigint, + service_name string, + "resource_attributes.service.namespace" string, + "resource_attributes.host.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', + 'greptime.semantic.entity.literal_resource.id' = 'resource_attributes.host.name'); + +insert into graph_trace_v1_peer values + ('2026-01-01 00:00:01', 'paired', 'server', 'client', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 50000000, 'backend', 'shop', 'node-v1'); + +select entity_type, entity_id, scope +from greptime_private.semantic_entities +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + +select src_type, src_id, dst_type, dst_id, rel_type, request_count, error_count +from greptime_private.semantic_relationships +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, src_type, src_id, dst_id; + +-- Empty components in the V1-compatible comma-separated list remain invalid. +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.host.id,'; +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.host.id,,service_name'; + +-- ALTER accepts JSON2 references; a missing key simply yields no entity. +alter table graph_trace_v2 set 'greptime.semantic.entity.custom.id' = 'scope_attributes.missing'; +alter table graph_trace_v2 set 'greptime.semantic.entity.host.descriptive' = 'resource_attributes.host.name'; +-- Preserve implicit host identity by explicitly declaring it as well. +alter table graph_trace_v2 set 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'; + +admin flush_table('graph_trace_v2'); +admin flush_table('graph_trace_v1_peer'); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, scope, entity_id_attrs, descriptive +from greptime_private.semantic_entities +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + +select src_type, src_id, dst_type, dst_id, rel_type, request_count, error_count +from greptime_private.semantic_relationships +where observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, src_type, src_id, dst_id; + +-- Quotes, backslashes and dots remain a literal key; numbers render as strings. +alter table graph_trace_v2 set 'greptime.semantic.entity.custom.id' = 'scope_attributes.a"b\c.d'; +select entity_type, entity_id +from greptime_private.semantic_entities +where entity_type = 'custom' + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00'; + +-- Path-looking suffixes are literal keys, not nested traversal or array access. +-- These missing keys are valid declarations and produce no entities. +alter table graph_trace_v2 set 'greptime.semantic.entity.literal_missing.id' = 'resource_attributes.host..name'; +alter table graph_trace_v2 set 'greptime.semantic.entity.literal_index.id' = 'resource_attributes.host[0]'; + +-- The root must exist and be one of the three JSON2 attribute columns. +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'span_events.name'; +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes.'; +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'missing.key'; +alter table graph_trace_v2 set 'greptime.semantic.entity.invalid.id' = 'resource_attributes'; + +create table graph_bad_v2 ( + ts timestamp time index, resource_attributes string +) with ('table_data_model' = 'greptime_trace_v2', + 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'); + +-- Ordinary JSON2 tables do not acquire Trace V2 attribute-reference semantics. +create table graph_not_trace ( + ts timestamp time index, resource_attributes json2 +) with ('append_mode' = 'true', + 'greptime.semantic.entity.host.id' = 'resource_attributes.host.id'); + +-- Physical convention columns win over conflicting JSON2 keys. +alter table graph_trace_v2 add column "resource_attributes.host.id" string; +alter table graph_trace_v2 add column "resource_attributes.service.namespace" string; +alter table graph_trace_v2 add column "span_attributes.db.namespace" string; +-- Restore the implicit host declaration for the physical-column check. +alter table graph_trace_v2 unset 'greptime.semantic.entity.host.id'; +alter table graph_trace_v2 unset 'greptime.semantic.entity.host.descriptive'; +insert into graph_trace_v2 + ("timestamp", duration_nano, trace_id, span_id, span_kind, span_status_code, + service_name, resource_attributes, span_attributes, + "resource_attributes.host.id", "resource_attributes.service.namespace", "span_attributes.db.namespace") +values ('2026-01-01 00:00:04', 100, 'physical', 'physical-client', 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', + 'physical-service', '{"host.id":"json-host","service.namespace":"json-namespace"}', + '{"db.namespace":"json-database","db":{"namespace":"nested-database"}}', + 'physical-host', 'physical-namespace', 'physical-database'); + +select entity_type, entity_id +from greptime_private.semantic_entities +where entity_type in ('host', 'service') + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by entity_type, entity_id; + +select src_id, dst_id, rel_type +from greptime_private.semantic_relationships +where src_id = 'physical-namespace/physical-service' + and observed_at >= '2026-01-01 00:00:00' and observed_at < '2026-01-01 00:01:00' +order by rel_type, dst_id; + +drop table graph_trace_v2; +drop table graph_trace_v1_peer; From b5199bc59a199875187d808158357a33b09e5f60 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:03:23 +0000 Subject: [PATCH 46/63] feat: add experimental Jev SQL filtering (#9265) * feat: add experimental Jev SQL filtering Signed-off-by: Lei, HUANG * feat: gate Jev filtering behind an opt-in Cargo feature Signed-off-by: Lei, HUANG * ci: verify Jev registration with default features Run the existing registry regression without the jev feature in both PR tests and merge-queue coverage, alongside the existing feature-enabled test runs. Signed-off-by: Lei, HUANG * docs: clarify Jev concurrency scope and stabilization work Document the per-expression/batch concurrency bound and track process-wide limiting, rate-limit backoff, and request budgets and metrics as stabilization prerequisites. Signed-off-by: Lei, HUANG * feat: enable renamed ai-functions feature by default Rename the Jev Cargo feature across the command, query, and function crates and enable it in their defaults. Update CI and documentation, retaining an isolated no-default-features registry check and the runtime API opt-in. Signed-off-by: Lei, HUANG * ci: remove extra AI feature-off checks Use the regular AI-enabled unit and coverage runs for the default feature configuration. Keep feature-off validation available locally and update the usage guide to match. Signed-off-by: Lei, HUANG * refactor: rename AI feature to ai_functions Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- .github/workflows/integration.yml | 2 +- .github/workflows/rust.yml | 4 +- Cargo.lock | 2 + docs/how-to/jev.md | 119 +++++++++ src/cmd/Cargo.toml | 2 + src/common/function/Cargo.toml | 10 +- src/common/function/src/function_registry.rs | 12 + src/common/function/src/scalars.rs | 2 + src/common/function/src/scalars/jev.rs | 209 +++++++++++++++ src/common/function/src/scalars/jev/tests.rs | 250 ++++++++++++++++++ src/query/Cargo.toml | 2 + .../standalone/common/function/jev.result | 66 +++++ .../cases/standalone/common/function/jev.sql | 32 +++ tests/runner/src/env/bare.rs | 2 +- 14 files changed, 709 insertions(+), 5 deletions(-) create mode 100644 docs/how-to/jev.md create mode 100644 src/common/function/src/scalars/jev.rs create mode 100644 src/common/function/src/scalars/jev/tests.rs create mode 100644 tests/cases/standalone/common/function/jev.result create mode 100644 tests/cases/standalone/common/function/jev.sql diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 2cbfa1e1541..c81dba199e9 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -81,7 +81,7 @@ jobs: # `cargo gc` will invoke `cargo build` with specified args run: | started_epoch="$(date +%s)" - cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index" + cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions" elapsed_secs=$(( $(date +%s) - started_epoch )) printf -- '- Binary build: **%ss**\n' "${elapsed_secs}" >>"${GITHUB_STEP_SUMMARY}" - name: Pack greptime binaries diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index b274931bc5b..1a093a2d2a5 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -200,7 +200,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 @@ -258,7 +258,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index + run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/Cargo.lock b/Cargo.lock index 7dadf276b7a..2c7f3333329 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2479,6 +2479,7 @@ dependencies = [ "arrow-cast 59.2.0", "arrow-schema 59.2.0", "async-trait", + "axum 0.8.4", "bincode", "catalog", "chrono", @@ -2518,6 +2519,7 @@ dependencies = [ "paste", "pretty_assertions", "regex", + "reqwest 0.13.4", "s2", "serde", "serde_json", diff --git a/docs/how-to/jev.md b/docs/how-to/jev.md new file mode 100644 index 00000000000..2a21f085946 --- /dev/null +++ b/docs/how-to/jev.md @@ -0,0 +1,119 @@ +# Natural-language filtering with Jev (experimental MVP) + +`jev(text, statement, threshold)` returns whether Jev judges the statement true +for the text with probability **greater than or equal to** the threshold. +It is an asynchronous SQL scalar function usable in `WHERE` and `SELECT`. + +## Start GreptimeDB + +Jev is compiled and registered by the **default-enabled Cargo feature +`ai_functions`**. Enable API evaluation separately in the server process: + +```sh +export GREPTIMEDB_EXPERIMENTAL_JEV=true +export JEV_API_KEY='' +cargo run -p cmd -- standalone start +``` + +If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. +The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to +`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). +The MVP uses environment variables rather than TOML configuration. +The Cargo feature includes the SQL function in the build; the runtime environment +variables enable and configure its API calls. Default compilation does not turn +on external API calls. To enable the feature explicitly, use +`--features ai_functions`. + +## Query + +For a new example table, quote `service` and `message` in the DDL because the +current parser treats them as keywords: + +```sql +CREATE TABLE events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); +``` + +```sql +SELECT occurred_at, service, message +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev( + (message), + 'The event reports that a payment still failed after retries.', + 0.8 + ) +ORDER BY occurred_at; +``` + +The first two arguments are strings and the third is a number in `[0, 1]`. +`(message)` is an ordinary parenthesized string expression. For several columns, +combine them explicitly, for example `concat(service, ': ', message)`. +Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. + +## MVP behavior + +- Each non-null row makes one HTTP request: the text is `state`, and the statement + is a `noul` question's `instructions`. The returned `noul` probability is compared + with the threshold locally; this is not the separate Choice/Score confidence. +- Up to eight requests run concurrently per expression/batch invocation, not per + query or process. Concurrent partitions and queries can exceed eight requests + in total. Each request has a 30-second timeout. Errors (including rate limits + and invalid responses) fail the query. There is no automatic retry or + cross-query cache. +- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the + candidate set, but SQL does not guarantee left-to-right predicate evaluation or + that `LIMIT` bounds the number of API calls. +- Text passed to the function is sent to the configured TypeSafe endpoint. + The API key remains in the server environment, not in SQL. +- This MVP is intended for standalone use. A distributed deployment would need + the environment on every node evaluating the function and separate validation. + +API reference: + +## Before stabilizing + +The experimental MVP still needs the following controls before stabilization: + +- A process-wide concurrency limit shared by Jev invocations, such as a semaphore. +- Bounded retries with exponential backoff for HTTP `429` and `529`, following + TypeSafe's rate-limit guidance. +- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. + +These are follow-up work, not guarantees provided by the current implementation. + +## Validation + +Check both the default build (where `jev` is registered) and an isolated +`common-function` build without default features. Jev's regular tests use a local +HTTP server and need no API key: + +```sh +cargo nextest run -p common-function +cargo nextest run -p common-function --no-default-features +cargo nextest run -p common-function --no-default-features --features ai_functions +``` + +The sqlness runner explicitly includes `ai_functions` when building its test binary: + +```sh +cargo sqlness bare -t jev +``` + +If using `--bins-dir`, provide a binary built with `ai_functions` (included by +default). CI covers the AI-enabled path through the regular unit tests and +sqlness tests; the feature-off checks above can be run locally when needed. + +An opt-in test calls the real service on synthetic payment events: + +```sh +source ~/.zshrc +GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ + -E 'test(test_jev_live)' --run-ignored ignored-only +``` diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 722e48ae90a..fff268dea45 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -21,11 +21,13 @@ required-features = ["dev-tools"] [features] default = [ + "ai_functions", "servers/pprof", "servers/mem-prof", "meta-srv/pg_kvbackend", "meta-srv/mysql_kvbackend", ] +ai_functions = ["query/ai_functions"] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] # Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index a1ea31bc32d..ee6cfac9dfb 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -10,8 +10,9 @@ workspace = true [features] testing = [] enterprise = [] -default = ["geo"] +default = ["geo", "ai_functions"] geo = ["geohash", "h3o", "s2", "wkt", "geo-types", "dep:geo"] +ai_functions = ["dep:futures", "dep:reqwest"] [dependencies] ahash.workspace = true @@ -43,6 +44,7 @@ datafusion-pg-catalog.workspace = true datafusion-physical-expr.workspace = true datatypes.workspace = true derive_more.workspace = true +futures = { workspace = true, optional = true } geo = { version = "0.29", optional = true } geo-types = { version = "0.7", optional = true } geohash = { version = "0.13", optional = true } @@ -57,6 +59,7 @@ num = "0.4" num-traits = "0.2" paste.workspace = true regex.workspace = true +reqwest = { workspace = true, optional = true } s2 = { version = "0.0.12", optional = true } serde.workspace = true serde_json.workspace = true @@ -71,6 +74,7 @@ wkt = { version = "0.11", optional = true } [dev-dependencies] approx = "0.5" +axum.workspace = true criterion.workspace = true futures.workspace = true pretty_assertions.workspace = true @@ -80,3 +84,7 @@ tokio.workspace = true [[bench]] name = "uddsketch" harness = false + +[package.metadata.cargo-udeps.ignore] +# Used only by the feature-gated Jev HTTP tests. +development = ["axum"] diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index c8e9874abc9..037ca5b175a 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -34,6 +34,8 @@ use crate::scalars::date::DateFunction; use crate::scalars::expression::ExpressionFunction; use crate::scalars::hll_count::HllCalcFunction; use crate::scalars::ip::IpFunctions; +#[cfg(feature = "ai_functions")] +use crate::scalars::jev::JevFunction; use crate::scalars::json::JsonFunction; use crate::scalars::matches::MatchesFunction; use crate::scalars::matches_term::MatchesTermFunction; @@ -226,6 +228,8 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| // Full text search function MatchesFunction::register(&function_registry); MatchesTermFunction::register(&function_registry); + #[cfg(feature = "ai_functions")] + JevFunction::register(&function_registry); // System and administration functions SystemFunction::register(&function_registry); @@ -379,6 +383,14 @@ mod tests { assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some()); } + #[test] + fn test_jev_registration_matches_feature() { + assert_eq!( + FUNCTION_REGISTRY.get_function("jev").is_some(), + cfg!(feature = "ai_functions") + ); + } + #[test] fn test_register_if_absent_registers_new_function() { let registry = FunctionRegistry::default(); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 306d1cd17e9..61fb7ec632c 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -17,6 +17,8 @@ pub(crate) mod date; pub mod expression; #[cfg(feature = "geo")] pub mod geo; +#[cfg(feature = "ai_functions")] +pub(crate) mod jev; pub mod json; pub mod matches; pub mod matches_term; diff --git a/src/common/function/src/scalars/jev.rs b/src/common/function/src/scalars/jev.rs new file mode 100644 index 00000000000..1c6475e6fac --- /dev/null +++ b/src/common/function/src/scalars/jev.rs @@ -0,0 +1,209 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::BooleanArray; +use arrow::datatypes::DataType; +use async_trait::async_trait; +use datafusion_common::cast::{as_float64_array, as_string_view_array}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err}; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use futures::{StreamExt, TryStreamExt, stream}; +use reqwest::Client; +use serde_json::{Value, json}; + +use crate::function_factory::ScalarFunctionFactory; +use crate::function_registry::FunctionRegistry; + +/// Experimental natural-language predicate: `jev(text, statement, threshold)`. +/// Each non-null row is a Noul question; its probability is compared with `>=`. +#[derive(PartialEq, Eq, Hash)] +pub(crate) struct JevFunction { + signature: Signature, + enabled: bool, + api_key: Option, + endpoint: String, + model: String, +} + +impl JevFunction { + pub(crate) fn register(registry: &FunctionRegistry) { + registry.register(ScalarFunctionFactory { + name: "jev".to_string(), + factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), + }); + } + + async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result { + let response: Value = client + .post(&self.endpoint) + .json(&json!({ + "model": self.model, + "state": text, + "questions": { + "matches": { "type": "noul", "instructions": statement } + } + })) + .send() + .await + .and_then(|response| response.error_for_status()) + .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))? + .json() + .await + .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?; + + let answer = &response["answers"]["matches"]; + let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p)); + match (answer["type"].as_str(), probability) { + (Some("noul"), Some(probability)) => Ok(probability), + _ => exec_err!( + "jev response must contain answers.matches with type noul and a probability in [0, 1]" + ), + } + } + + fn client(&self) -> Result { + if !self.enabled { + return exec_err!( + "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it" + ); + } + let key = self + .api_key + .as_deref() + .filter(|key| !key.trim().is_empty()) + .ok_or_else(|| { + exec_datafusion_err!("jev requires the JEV_API_KEY environment variable") + })?; + let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) + .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; + authorization.set_sensitive(true); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::AUTHORIZATION, authorization); + Client::builder() + .default_headers(headers) + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}")) + } +} + +impl Default for JevFunction { + fn default() -> Self { + Self { + // External model evaluations must not be constant-folded during planning. + signature: Signature::exact( + vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64], + Volatility::Volatile, + ), + enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), + api_key: std::env::var("JEV_API_KEY").ok(), + endpoint: std::env::var("JEV_ENDPOINT") + .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), + model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), + } + } +} + +impl fmt::Debug for JevFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Credentials must never appear in plans or diagnostic output. + f.debug_struct("JevFunction") + .field("signature", &self.signature) + .field("enabled", &self.enabled) + .field("model", &self.model) + .finish_non_exhaustive() + } +} + +impl ScalarUDFImpl for JevFunction { + fn name(&self) -> &str { + "jev" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("jev can only be called from async contexts") + } +} + +#[async_trait] +impl AsyncScalarUDFImpl for JevFunction { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = args + .args + .into_iter() + .map(|arg| arg.into_array(args.number_rows)) + .collect::>>()?; + let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?; + let texts = as_string_view_array(&texts)?; + let statements = as_string_view_array(&statements)?; + let thresholds = as_float64_array(&thresholds)?; + let rows: Vec<_> = texts + .iter() + .zip(statements.iter()) + .zip(thresholds.iter()) + .map(|((text, statement), threshold)| Some((text?, statement?, threshold?))) + .collect(); + + // Validate the whole batch before making any billable requests. + for (_, _, threshold) in rows.iter().flatten() { + if !(0.0..=1.0).contains(threshold) { + return exec_err!("jev threshold must be finite and in [0, 1]"); + } + } + if rows.iter().all(Option::is_none) { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + args.number_rows, + )))); + } + + // Reuse connections within the batch; bound in-flight requests and preserve row order. + let client = self.client()?; + let requests: Vec<_> = rows + .into_iter() + .map(|row| { + let client = &client; + async move { + match row { + Some((text, statement, threshold)) => self + .evaluate(client, text, statement) + .await + .map(|p| Some(p >= threshold)), + None => Ok(None), + } + } + }) + .collect(); + // This limit is per expression/batch invocation, not per query or process. + // Concurrent partitions and queries can each have their own in-flight requests. + let matches: Vec> = stream::iter(requests).buffered(8).try_collect().await?; + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches)))) + } +} diff --git a/src/common/function/src/scalars/jev/tests.rs b/src/common/function/src/scalars/jev/tests.rs new file mode 100644 index 00000000000..82d7ab03013 --- /dev/null +++ b/src/common/function/src/scalars/jev/tests.rs @@ -0,0 +1,250 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use axum::http::{HeaderMap, StatusCode}; +use axum::routing::post; +use axum::{Json, Router}; +use datafusion::assert_batches_eq; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +use super::*; +use crate::function::FunctionContext; +use crate::function_registry::FUNCTION_REGISTRY; + +struct MockServer { + endpoint: String, + task: JoinHandle<()>, +} + +impl MockServer { + async fn start(router: Router) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + Self { endpoint, task } + } + + fn function(&self) -> JevFunction { + JevFunction { + enabled: true, + api_key: Some("test-key".to_string()), + endpoint: self.endpoint.clone(), + model: "test-model".to_string(), + ..Default::default() + } + } +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn context(function: JevFunction) -> SessionContext { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); + ctx +} + +#[tokio::test] +async fn test_jev_sql_filters_using_noul_probability() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post( + |headers: HeaderMap, Json(request): Json| async move { + assert_eq!(headers["authorization"], "Bearer test-key"); + assert_eq!(request["model"], "test-model"); + assert_eq!(request["questions"]["matches"]["type"], "noul"); + // The SQL statement is passed through verbatim, not rewritten as a prompt. + assert_eq!( + request["questions"]["matches"]["instructions"], + "payment failed" + ); + let probability = match request["state"].as_str().unwrap() { + "failed" => { + // Complete this row after later rows to exercise result alignment. + tokio::time::sleep(Duration::from_millis(20)).await; + 0.95 + } + "boundary" => 0.8, + "recovered" => 0.1, + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) + }, + ), + )) + .await; + let ctx = context(server.function()); + let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)"; + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + ["+----+", "| id |", "+----+", "| 1 |", "| 2 |", "+----+"], + &batches + ); + + let batches = ctx + .sql(&format!( + "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+---------+", + "| id | matched |", + "+----+---------+", + "| 1 | true |", + "| 2 | true |", + "| 3 | false |", + "| 4 | |", + "+----+---------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() { + let ctx = context(JevFunction { + enabled: true, + api_key: None, + ..Default::default() + }); + let batches = ctx + .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c") + .await.unwrap().collect().await.unwrap(); + assert_batches_eq!( + [ + "+---+---+---+", + "| a | b | c |", + "+---+---+---+", + "| | | |", + "+---+---+---+" + ], + &batches + ); + + for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] { + let error = ctx + .sql(&format!("SELECT jev('text', 'condition', {threshold})")) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("jev threshold must be finite and in [0, 1]"), + "{error}" + ); + } + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); + + let ctx = context(JevFunction { + enabled: false, + api_key: None, + ..Default::default() + }); + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), + "{error}" + ); +} + +#[tokio::test] +async fn test_jev_bad_api_responses_fail_the_query() { + for (status, body, expected) in [ + (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"), + (StatusCode::OK, "not json", "jev response is not valid JSON"), + ( + StatusCode::OK, + r#"{"answers":{}}"#, + "jev response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, + "jev response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, + "jev response must contain", + ), + ] { + let server = MockServer::start( + Router::new().route("/v1/systemone", post(move || async move { (status, body) })), + ) + .await; + let ctx = context(server.function()); + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + } +} + +#[tokio::test] +#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] +async fn test_jev_live() { + let ctx = SessionContext::new(); + let function = FUNCTION_REGISTRY.get_function("jev").unwrap(); + ctx.register_udf(function.provide(FunctionContext::default())); + let batches = ctx.sql( + "SELECT id FROM (VALUES + (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), + (2, 'Payment succeeded on the second retry; the payment is complete.'), + (3, 'User logged in successfully.') + ) AS events(id, message) + WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8) + ORDER BY id", + ).await.unwrap().collect().await.unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); +} diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index b6dd2d1fdf4..480c95a5adb 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -8,6 +8,8 @@ license.workspace = true workspace = true [features] +default = ["ai_functions"] +ai_functions = ["common-function/ai_functions"] enterprise = ["sql/enterprise"] vector_index = [] diff --git a/tests/cases/standalone/common/function/jev.result b/tests/cases/standalone/common/function/jev.result new file mode 100644 index 00000000000..4f6d91d2867 --- /dev/null +++ b/tests/cases/standalone/common/function/jev.result @@ -0,0 +1,66 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +SELECT jev('message', NULL, 0.8) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE jev_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +Affected Rows: 0 + +INSERT INTO jev_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +Affected Rows: 2 + +SELECT occurred_at, service, message FROM jev_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) +ORDER BY occurred_at; + ++-------------+---------+---------+ +| occurred_at | service | message | ++-------------+---------+---------+ ++-------------+---------+---------+ + +DROP TABLE jev_events; + +Affected Rows: 0 + +-- Invalid thresholds fail locally before any HTTP request. +SELECT jev('message', 'condition', -0.1); + +Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] + +SELECT jev('message', 'condition', 1.1); + +Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] + diff --git a/tests/cases/standalone/common/function/jev.sql b/tests/cases/standalone/common/function/jev.sql new file mode 100644 index 00000000000..f4c252c9273 --- /dev/null +++ b/tests/cases/standalone/common/function/jev.sql @@ -0,0 +1,32 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; + +SELECT jev('message', NULL, 0.8) AS matched; + +SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE jev_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +INSERT INTO jev_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +SELECT occurred_at, service, message FROM jev_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) +ORDER BY occurred_at; + +DROP TABLE jev_events; + +-- Invalid thresholds fail locally before any HTTP request. +SELECT jev('message', 'condition', -0.1); + +SELECT jev('message', 'condition', 1.1); diff --git a/tests/runner/src/env/bare.rs b/tests/runner/src/env/bare.rs index 8fc3855edbb..dccd3673b16 100644 --- a/tests/runner/src/env/bare.rs +++ b/tests/runner/src/env/bare.rs @@ -694,7 +694,7 @@ impl Env { "--bin", "greptime", "--features", - "pg_kvbackend,mysql_kvbackend,vector_index", + "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions", ]) .output() .expect("Failed to start GreptimeDB"); From 2a63ea2660991895168e5b43f8f88bd102523bdf Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 04:01:31 +0000 Subject: [PATCH 47/63] feat(log-store): add object store WAL reads and obsolete watermarks (#9248) * feat(log-store): add object store WAL reads and obsolete watermarks Signed-off-by: jeremyhi * perf(log-store): reuse decoded WAL payload allocations Signed-off-by: jeremyhi * refactor(log-store): clarify WAL region mismatch error Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/error.rs | 18 +- src/log-store/src/object_store_wal/store.rs | 488 +++++++++++++++++--- 2 files changed, 450 insertions(+), 56 deletions(-) diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 6826be716f3..de304e192d2 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -324,6 +324,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Object store WAL region mismatch, supplied region: {}, {}", + region_id, + reason + ))] + MismatchedWalRegion { + region_id: RegionId, + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "Invalid WAL entry range, region: {}, start: {}, end: {}", region_id, @@ -400,9 +412,8 @@ pub enum Error { location: Location, }, - /// The store supports construction, recovery, stopping and latest-entry queries; - /// other log-store operations return this error. - #[snafu(display("Object store WAL operation is not supported yet"))] + /// Appending to the object store WAL is unsupported. + #[snafu(display("Object store WAL operation is not supported"))] UnsupportedObjectStoreWalOperation { #[snafu(implicit)] location: Location, @@ -440,6 +451,7 @@ impl ErrorExt for Error { | OverrideCompactedEntry { .. } | InvalidWalObjectStore { .. } | MismatchedWalPrefix { .. } + | MismatchedWalRegion { .. } | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, StartWalTask { .. } | StopWalTask { .. } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index e578bc79126..c0d9cc89363 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -12,21 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Construction and footer-based recovery of an object store WAL. +//! Object store WAL construction, recovery and region reads. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::fmt; use std::ops::Range; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError, RwLock}; use std::time::Duration; +use async_stream::try_stream; use bytes::Bytes; use common_wal::config::object_store::ObjectStoreWalConfig; use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStore; use snafu::{IntoError, OptionExt, ResultExt, ensure}; -use store_api::logstore::entry::Entry; +use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; use store_api::storage::RegionId; @@ -34,13 +35,13 @@ use tokio::sync::{mpsc, oneshot}; use crate::error::{ CorruptedWalObjectSnafu, Error, InvalidProviderSnafu, InvalidWalObjectSnafu, - InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, ObjectStoreWalSnafu, Result, - UnsupportedObjectStoreWalOperationSnafu, + InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, MismatchedWalRegionSnafu, + ObjectStoreWalSnafu, Result, UnsupportedObjectStoreWalOperationSnafu, }; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ FixedTrailer, FooterEntry, HEADER_LEN, MIN_OBJECT_LEN, TRAILER_LEN, decode_footer, - decode_header, decode_trailer, footer_range, verify_segment_ranges, + decode_header, decode_segment, decode_trailer, footer_range, verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; @@ -58,11 +59,14 @@ pub(crate) struct ObjectStoreLogStore { prefix: String, io: Arc, catalog: Arc>, + obsolete_entry_ids: ObsoleteEntryIds, terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, } +type ObsoleteEntryIds = Arc>>; + type TerminalError = Arc>>>; impl fmt::Debug for ObjectStoreLogStore { @@ -122,6 +126,7 @@ impl ObjectStoreLogStore { prefix, io, catalog, + obsolete_entry_ids: ObsoleteEntryIds::default(), terminal_error, stopped, command_tx, @@ -152,6 +157,19 @@ impl ObjectStoreLogStore { Ok(provider.region_id) } + fn check_region(&self, provider: &Provider, region_id: RegionId) -> Result<()> { + self.check_terminal()?; + let provider_region = self.region_of(provider)?; + ensure!( + provider_region == region_id, + MismatchedWalRegionSnafu { + region_id, + reason: format!("provider belongs to region {provider_region}"), + } + ); + Ok(()) + } + fn check_terminal(&self) -> Result<()> { match terminal(&self.terminal_error) { Some(error) => Err(shared(&error)), @@ -160,6 +178,16 @@ impl ObjectStoreLogStore { } } +/// Records the obsolete watermark of `region_id`, which never moves down. +fn record_obsolete(obsolete_entry_ids: &ObsoleteEntryIds, region_id: RegionId, entry_id: EntryId) { + obsolete_entry_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(region_id) + .and_modify(|current| *current = (*current).max(entry_id)) + .or_insert(entry_id); +} + fn positive_bytes(bytes: u64, name: &str) -> Result { usize::try_from(bytes) .ok() @@ -193,48 +221,127 @@ impl LogStore for ObjectStoreLogStore { UnsupportedObjectStoreWalOperationSnafu.fail() } + /// Reads the provider's region from `entry_id`, hiding obsolete entries. + /// The catalog locates each segment, so a caller-supplied index is unnecessary. async fn read( &self, - _provider: &Provider, - _entry_id: EntryId, + provider: &Provider, + entry_id: EntryId, _index: Option, ) -> Result> { - UnsupportedObjectStoreWalOperationSnafu.fail() + self.check_terminal()?; + let region_id = self.region_of(provider)?; + let obsolete = self + .obsolete_entry_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(®ion_id) + .copied(); + let start_entry_id = + entry_id.max(obsolete.map_or(0, |obsolete| obsolete.saturating_add(1))); + let objects = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + match catalog.region_max_entry_id(region_id) { + Some(max_entry_id) + if start_entry_id <= max_entry_id && obsolete != Some(EntryId::MAX) => + { + catalog + .objects_for_entry_range(region_id, start_entry_id, max_entry_id)? + .into_iter() + .map(|(object_seq, entry)| (object_seq, entry.clone())) + .collect::>() + } + _ => Vec::new(), + } + }; + + let io = self.io.clone(); + let provider = provider.clone(); + Ok(Box::pin(try_stream! { + for (object_seq, footer_entry) in objects { + let bytes = io + .get_range(object_seq, footer_entry.segment_offset, footer_entry.segment_len) + .await?; + let records = decode_segment(&bytes, &footer_entry) + .with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + })?; + let entries = records + .into_iter() + .filter(|record| record.entry_id >= start_entry_id) + .map(|record| { + Entry::Naive(NaiveEntry { + provider: provider.clone(), + region_id, + entry_id: record.entry_id, + data: record.payload.into(), + }) + }) + .collect::>(); + if !entries.is_empty() { + yield entries; + } + } + })) } - async fn create_namespace(&self, _ns: &Provider) -> Result<()> { - UnsupportedObjectStoreWalOperationSnafu.fail() + async fn create_namespace(&self, ns: &Provider) -> Result<()> { + self.check_terminal()?; + self.region_of(ns).map(|_| ()) } - async fn delete_namespace(&self, _ns: &Provider) -> Result<()> { - UnsupportedObjectStoreWalOperationSnafu.fail() + async fn delete_namespace(&self, ns: &Provider) -> Result<()> { + self.check_terminal()?; + self.region_of(ns).map(|_| ()) } async fn list_namespaces(&self) -> Result> { - UnsupportedObjectStoreWalOperationSnafu.fail() + self.check_terminal()?; + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + let regions = catalog + .objects_in_order() + .flat_map(|(_, footer)| footer.iter().map(|entry| entry.region_id)) + .collect::>(); + Ok(regions + .into_iter() + .map(|region_id| Provider::object_store_provider(region_id, self.prefix.clone())) + .collect()) } + /// Records a memory-only watermark. The caller re-establishes it after a restart. async fn obsolete( &self, - _provider: &Provider, - _region_id: RegionId, - _entry_id: EntryId, + provider: &Provider, + region_id: RegionId, + entry_id: EntryId, ) -> Result<()> { - UnsupportedObjectStoreWalOperationSnafu.fail() + self.check_region(provider, region_id)?; + record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + Ok(()) } - async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { - UnsupportedObjectStoreWalOperationSnafu.fail() + /// Hides every entry of the region in memory. The caller re-establishes + /// the watermark after a restart. + async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { + self.check_region(provider, region_id)?; + record_obsolete(&self.obsolete_entry_ids, region_id, EntryId::MAX); + Ok(()) } fn entry( &self, - _data: Vec, - _entry_id: EntryId, - _region_id: RegionId, - _provider: &Provider, + data: Vec, + entry_id: EntryId, + region_id: RegionId, + provider: &Provider, ) -> Result { - UnsupportedObjectStoreWalOperationSnafu.fail() + self.check_region(provider, region_id)?; + Ok(Entry::Naive(NaiveEntry { + provider: provider.clone(), + region_id, + entry_id, + data, + })) } fn latest_entry_id(&self, provider: &Provider) -> Result { @@ -1058,15 +1165,32 @@ mod tests { #[tokio::test] async fn test_store_rejects_foreign_providers() { let store = open(memory_store(), &eager()).await; - let other = Provider::object_store_provider(region(1), "other/prefix".to_string()); - assert!(matches!(store.latest_entry_id(&other).unwrap_err(), - Error::MismatchedWalPrefix { expected, actual, .. } - if expected == PREFIX && actual == "other/prefix")); - let raft = Provider::raft_engine_provider(region(1).as_u64()); - assert!(matches!( - store.latest_entry_id(&raft), - Err(Error::InvalidProvider { .. }) - )); + let region_id = region(1); + let other = Provider::object_store_provider(region_id, "other/prefix".to_string()); + let raft = Provider::raft_engine_provider(region_id.as_u64()); + for foreign in [&other, &raft] { + let errors = [ + store.latest_entry_id(foreign).unwrap_err(), + store.read(foreign, 0, None).await.err().unwrap(), + store.create_namespace(foreign).await.unwrap_err(), + store.delete_namespace(foreign).await.unwrap_err(), + store.obsolete(foreign, region_id, 1).await.unwrap_err(), + store.obsolete_all(foreign, region_id).await.unwrap_err(), + store.entry(Vec::new(), 1, region_id, foreign).unwrap_err(), + ]; + for error in errors { + if foreign == &raft { + assert!(matches!(error, Error::InvalidProvider { .. }), "{error:?}"); + } else { + assert!( + matches!(&error, + Error::MismatchedWalPrefix { expected, actual, .. } + if expected == PREFIX && actual == "other/prefix"), + "{error:?}" + ); + } + } + } store.stop().await.unwrap(); } @@ -1089,6 +1213,7 @@ mod tests { prefix: PREFIX.to_string(), io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), catalog: Arc::default(), + obsolete_entry_ids: ObsoleteEntryIds::default(), terminal_error: Arc::default(), stopped: Arc::new(AtomicBool::new(false)), command_tx, @@ -1125,10 +1250,27 @@ mod tests { CorruptedWalObjectSnafu { reason: "second" }.build(), ); assert!(Arc::ptr_eq(&first, &second)); - match store.latest_entry_id(&provider(region(1))).unwrap_err() { - Error::ObjectStoreWal { source, .. } => assert!(Arc::ptr_eq(&first, &source)), - error => panic!("{error:?}"), + let region_id = region(1); + let provider = provider(region_id); + let errors = [ + store.latest_entry_id(&provider).unwrap_err(), + store.read(&provider, 0, None).await.err().unwrap(), + store.create_namespace(&provider).await.unwrap_err(), + store.delete_namespace(&provider).await.unwrap_err(), + store.list_namespaces().await.unwrap_err(), + store.obsolete(&provider, region_id, 1).await.unwrap_err(), + store.obsolete_all(&provider, region_id).await.unwrap_err(), + store + .entry(Vec::new(), 1, region_id, &provider) + .unwrap_err(), + ]; + for error in errors { + match error { + Error::ObjectStoreWal { source, .. } => assert!(Arc::ptr_eq(&first, &source)), + error => panic!("{error:?}"), + } } + assert!(store.obsolete_entry_ids.lock().unwrap().is_empty()); store.stop().await.unwrap(); } @@ -1267,30 +1409,270 @@ mod tests { } #[tokio::test] - async fn test_store_unsupported_operations() { + async fn test_store_append_is_unsupported() { let store = open(memory_store(), &eager()).await; - let region = region(1); - let provider = provider(region); + let error = store.append_batch(Vec::new()).await.unwrap_err(); + assert!(matches!( + error, + Error::UnsupportedObjectStoreWalOperation { .. } + )); + assert_eq!( + common_error::status_code::StatusCode::Unsupported, + error.status_code() + ); + store.stop().await.unwrap(); + } + + async fn read_entries( + store: &ObjectStoreLogStore, + region_id: RegionId, + start: EntryId, + ) -> Vec { + store + .read(&provider(region_id), start, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect() + } + + fn expected_entries(region_id: RegionId, entries: &[(EntryId, &str)]) -> Vec { + entries + .iter() + .map(|(entry_id, data)| { + Entry::Naive(NaiveEntry { + provider: provider(region_id), + region_id, + entry_id: *entry_id, + data: data.as_bytes().to_vec(), + }) + }) + .collect() + } + + #[tokio::test] + async fn test_store_reads_only_region_segments_in_entry_order() { + let object_store = memory_store(); + let mut expected_ranges = HashMap::<_, Vec<_>>::new(); + for (seq, first) in [(1, 10), (3, 30), (7, 70)] { + let records = [region(2), region(1)] + .into_iter() + .flat_map(|region_id| { + (first..first + 3).map(move |entry_id| Record { + region_id, + entry_id, + payload: Bytes::from(format!("r{}-e{entry_id}", region_id.region_number())), + }) + }) + .collect::>(); + put_records(&object_store, seq, &records).await; + let bytes = object_store + .read(&object_path(&object_store, seq)) + .await + .unwrap() + .to_vec(); + for entry in footer_of(&bytes).1 { + expected_ranges.entry(entry.region_id).or_default().push(( + seq, + entry.segment_offset, + entry.segment_len, + )); + } + } + let (io, reads) = RecordingIo::over(object_store); + let store = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap(); + for number in [1, 2] { + reads.lock().unwrap().clear(); + let entries = read_entries(&store, region(number), 11).await; + let ids = [11, 12, 30, 31, 32, 70, 71, 72]; + let expected = ids + .into_iter() + .map(|entry_id| { + Entry::Naive(NaiveEntry { + provider: provider(region(number)), + region_id: region(number), + entry_id, + data: format!("r{number}-e{entry_id}").into_bytes(), + }) + }) + .collect::>(); + assert_eq!(expected, entries); + assert_eq!(expected_ranges[®ion(number)], *reads.lock().unwrap()); + } + reads.lock().unwrap().clear(); + assert_eq!( + Vec::::new(), + read_entries(&store, region(1), 73).await + ); + assert_eq!( + Vec::::new(), + read_entries(&store, region(3), 0).await + ); + assert!(reads.lock().unwrap().is_empty()); + assert_eq!( + vec![provider(region(1)), provider(region(2))], + store.list_namespaces().await.unwrap() + ); + store.create_namespace(&provider(region(3))).await.unwrap(); + store.delete_namespace(&provider(region(1))).await.unwrap(); + assert_eq!( + vec![provider(region(1)), provider(region(2))], + store.list_namespaces().await.unwrap() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_obsolete_hides_entries_from_read_only() { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[10, 11, 12]).await; + put_object(&object_store, 1, region(1), &[20]).await; + put_object(&object_store, 2, region(2), &[10]).await; + let store = open(object_store.clone(), &eager()).await; + let p = provider(region(1)); + let all = expected_entries( + region(1), + &[(10, "e10"), (11, "e11"), (12, "e12"), (20, "e20")], + ); + store.obsolete(&p, region(1), 9).await.unwrap(); + assert_eq!(all, read_entries(&store, region(1), 0).await); + store.obsolete(&p, region(1), 11).await.unwrap(); + let remaining = expected_entries(region(1), &[(12, "e12"), (20, "e20")]); + assert_eq!(remaining, read_entries(&store, region(1), 0).await); + assert_eq!( + expected_entries(region(1), &[(20, "e20")]), + read_entries(&store, region(1), 20).await + ); + assert_eq!(20, latest(&store, region(1))); + store.obsolete(&p, region(1), 10).await.unwrap(); + assert_eq!(remaining, read_entries(&store, region(1), 0).await); + store.obsolete_all(&p, region(1)).await.unwrap(); + store.obsolete(&p, region(1), 0).await.unwrap(); + assert_eq!( + Vec::::new(), + read_entries(&store, region(1), 0).await + ); + assert_eq!(20, latest(&store, region(1))); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 0).await + ); + store.stop().await.unwrap(); + let reopened = open(object_store, &eager()).await; + assert_eq!(all, read_entries(&reopened, region(1), 0).await); + reopened + .obsolete(&p, region(1), EntryId::MAX) + .await + .unwrap(); + assert_eq!( + Vec::::new(), + read_entries(&reopened, region(1), 0).await + ); + assert_eq!(20, latest(&reopened, region(1))); + reopened.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_entry_and_watermarks_reject_another_region() { + let store = open(memory_store(), &eager()).await; + let p = provider(region(2)); let errors = [ - store.append_batch(Vec::new()).await.unwrap_err(), - store.read(&provider, 0, None).await.err().unwrap(), - store.create_namespace(&provider).await.unwrap_err(), - store.delete_namespace(&provider).await.unwrap_err(), - store.list_namespaces().await.unwrap_err(), - store.obsolete(&provider, region, 1).await.unwrap_err(), - store.obsolete_all(&provider, region).await.unwrap_err(), - store.entry(Vec::new(), 1, region, &provider).unwrap_err(), + store + .entry(b"payload".to_vec(), 7, region(1), &p) + .unwrap_err(), + store.obsolete(&p, region(1), 7).await.unwrap_err(), + store.obsolete_all(&p, region(1)).await.unwrap_err(), ]; for error in errors { - assert!(matches!( - error, - Error::UnsupportedObjectStoreWalOperation { .. } - )); + assert!( + matches!(&error, Error::MismatchedWalRegion { region_id, reason, .. } + if *region_id == region(1) && reason == &format!("provider belongs to region {}", region(2))) + ); assert_eq!( - common_error::status_code::StatusCode::Unsupported, + common_error::status_code::StatusCode::InvalidArguments, error.status_code() ); } + assert!(store.obsolete_entry_ids.lock().unwrap().is_empty()); + assert_eq!( + expected_entries(region(2), &[(7, "payload")]), + vec![store.entry(b"payload".to_vec(), 7, region(2), &p).unwrap()] + ); + assert!(store.list_namespaces().await.unwrap().is_empty()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { + let object_store = memory_store(); + put_records( + &object_store, + 0, + &[ + Record { + region_id: region(1), + entry_id: 1, + payload: Bytes::from_static(b"a1"), + }, + Record { + region_id: region(2), + entry_id: 1, + payload: Bytes::from_static(b"b1"), + }, + ], + ) + .await; + put_object(&object_store, 1, region(2), &[10]).await; + let path = object_path(&object_store, 0); + let bytes = object_store.read(&path).await.unwrap().to_vec(); + let footer = footer_of(&bytes).1; + let corrupt = &footer[1]; + corrupt_object(&object_store, &path, |bytes| { + bytes[(corrupt.segment_offset + corrupt.segment_len - 1) as usize] ^= 1; + }) + .await; + let (io, reads) = RecordingIo::over(object_store); + let store = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap(); + assert_eq!(1, latest(&store, region(1))); + assert_eq!(10, latest(&store, region(2))); + reads.lock().unwrap().clear(); + assert_eq!( + expected_entries(region(1), &[(1, "a1")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!( + vec![(0, footer[0].segment_offset, footer[0].segment_len)], + *reads.lock().unwrap() + ); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 10).await + ); + reads.lock().unwrap().clear(); + let error = store + .read(&provider(region(2)), 0, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert_invalid_object( + &error, + &path, + &format!("segment of region {} checksum mismatch", region(2)), + ); + assert_eq!( + vec![(0, corrupt.segment_offset, corrupt.segment_len)], + *reads.lock().unwrap() + ); store.stop().await.unwrap(); } From e65f10c8adedfa188bf903071a3c49ed35605d17 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 04:02:40 +0000 Subject: [PATCH 48/63] fix(query): stop encoding oversized dynamic filters (#9267) * fix(query): stop encoding oversized dynamic filters Signed-off-by: Dennis Zhuang * test(query): stop bounded encoding partway through an IN list Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/query/src/request.rs | 137 ++++++++++++++++++++++++++++++-- 1 file changed, 130 insertions(+), 7 deletions(-) diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 56683818618..86c6758aa25 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -15,20 +15,26 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; +use std::cell::Cell; use std::sync::Arc; use api::v1::region::RegionRequestHeader; use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::{Column, InListExpr, lit}; -use datafusion::physical_plan::PhysicalExpr; use datafusion::physical_plan::joins::HashTableLookupExpr; +use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{DataFusionError, Result as DataFusionResult}; use datafusion_expr::LogicalPlan; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; use datafusion_proto::physical_plan::from_proto::parse_physical_expr; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use datafusion_proto::protobuf::PhysicalExprNode; +use datafusion_proto::physical_plan::to_proto::{ + serialize_physical_expr, serialize_physical_expr_with_converter, +}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use prost::Message; use serde::{Deserialize, Serialize}; use snafu::ensure; @@ -120,6 +126,69 @@ fn encode_physical_expr_to_bytes(expr: &Arc) -> DataFusionResu Ok(bytes) } +struct BoundedPhysicalExprEncoder { + max_payload_bytes: usize, + encoded_bytes: Cell, +} + +impl PhysicalProtoConverterExtension for BoundedPhysicalExprEncoder { + fn proto_to_execution_plan( + &self, + proto: &PhysicalPlanNode, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> DataFusionResult> { + DefaultPhysicalProtoConverter {}.proto_to_execution_plan(proto, ctx) + } + + fn execution_plan_to_proto( + &self, + plan: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> DataFusionResult { + DefaultPhysicalProtoConverter {}.execution_plan_to_proto(plan, codec) + } + + fn proto_to_physical_expr( + &self, + proto: &PhysicalExprNode, + input_schema: &datafusion::arrow::datatypes::Schema, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> DataFusionResult> { + DefaultPhysicalProtoConverter {}.proto_to_physical_expr(proto, input_schema, ctx) + } + + fn physical_expr_to_proto( + &self, + expr: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> DataFusionResult { + let before_children = self.encoded_bytes.get(); + let proto = serialize_physical_expr_with_converter(expr, codec, self)?; + // Replace the children's contribution with the whole subtree, including protobuf framing. + // Checking each child stops large IN lists before materializing all their Arrow IPC literals. + let encoded_bytes = before_children.saturating_add(proto.encoded_len()); + self.encoded_bytes.set(encoded_bytes); + validate_payload_size(encoded_bytes, self.max_payload_bytes) + .map_err(DataFusionError::from)?; + Ok(proto) + } +} + +fn encode_bounded_physical_expr_to_bytes( + expr: &Arc, + max_payload_bytes: usize, +) -> Result, CommonQueryError> { + let codec = DefaultPhysicalExtensionCodec {}; + let encoder = BoundedPhysicalExprEncoder { + max_payload_bytes, + encoded_bytes: Cell::new(0), + }; + let proto = encoder.physical_expr_to_proto(expr, &codec); + // Preserve the size error so the caller can retry with bounds only. + validate_payload_size(encoder.encoded_bytes.get(), max_payload_bytes)?; + Ok(proto.map_err(CommonQueryError::from)?.encode_to_vec()) +} + fn encode_remote_dyn_filter_expr( expr: &Arc, max_payload_bytes: usize, @@ -127,9 +196,7 @@ fn encode_remote_dyn_filter_expr( ) -> Result, CommonQueryError> { let expr = portable_remote_dyn_filter_expr(Arc::clone(expr), bounds_only) .map_err(CommonQueryError::from)?; - let bytes = encode_physical_expr_to_bytes(&expr).map_err(CommonQueryError::from)?; - validate_payload_size(bytes.len(), max_payload_bytes)?; - Ok(bytes) + encode_bounded_physical_expr_to_bytes(&expr, max_payload_bytes) } fn portable_remote_dyn_filter_expr( @@ -522,6 +589,62 @@ mod tests { assert!(decoded_display.contains("<=")); } + #[test] + fn bounded_encoding_preserves_bytes_and_exact_limit() { + let schema = Schema::new(vec![Field::new("device_id", DataType::Int32, false)]); + let device_id = Arc::new(Column::new("device_id", 0)) as Arc; + let in_list = Arc::new( + InListExpr::try_new( + Arc::clone(&device_id), + (0..64).map(lit).collect(), + false, + &schema, + ) + .unwrap(), + ) as Arc; + let lower_bound = Arc::new(BinaryExpr::new(device_id, Operator::GtEq, lit(10i32))) + as Arc; + let expr = + Arc::new(BinaryExpr::new(lower_bound, Operator::And, in_list)) as Arc; + let expected = serialize_physical_expr(&expr, &DefaultPhysicalExtensionCodec {}) + .unwrap() + .encode_to_vec(); + + assert_eq!( + encode_bounded_physical_expr_to_bytes(&expr, expected.len()).unwrap(), + expected + ); + assert!(matches!( + encode_bounded_physical_expr_to_bytes(&expr, expected.len() - 1), + Err(CommonQueryError::DynFilterPayloadTooLarge { .. }) + )); + } + + #[test] + fn bounded_encoding_stops_before_later_children() { + let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]); + let expr = Arc::new( + InListExpr::try_new( + Arc::new(Column::new("host", 0)), + (0..8).map(|i| lit(i.to_string().repeat(1024))).collect(), + false, + &schema, + ) + .unwrap(), + ) as Arc; + let full_size = encode_physical_expr_to_bytes(&expr).unwrap().len(); + assert!(full_size > 8192); + + let error = encode_bounded_physical_expr_to_bytes(&expr, 2048).unwrap_err(); + let CommonQueryError::DynFilterPayloadTooLarge { + payload_size_bytes, .. + } = error + else { + panic!("expected payload size error, got {error:?}"); + }; + assert!((2048..4096).contains(&payload_size_bytes)); + } + #[test] fn dyn_filter_payload_rejects_oversized_payload() { let expr: Arc = Arc::new(Column::new("host", 0)); From 67497f14e57880a3dd98046991184aa07ae8f0f4 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 22 Sep 2026 04:09:59 +0000 Subject: [PATCH 49/63] refactor: isolate logical table preparation and reuse Flow notifications (#9212) refactor: isolate logical table preparation and share Flow notifications Signed-off-by: WenyXu --- src/servers/src/batcher/logical_table.rs | 945 +----------------- .../src/batcher/logical_table/batch.rs | 569 ++++++----- .../batcher/logical_table/batch_convert.rs | 4 +- .../batcher/logical_table/flow_notifier.rs | 326 ++++++ .../batcher/logical_table/pending_worker.rs | 132 +-- .../src/batcher/logical_table/region_write.rs | 148 +-- .../src/batcher/logical_table/tables.rs | 551 ++++++++++ .../src/batcher/logical_table/test_util.rs | 25 - 8 files changed, 1330 insertions(+), 1370 deletions(-) create mode 100644 src/servers/src/batcher/logical_table/flow_notifier.rs create mode 100644 src/servers/src/batcher/logical_table/tables.rs diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index 8e7bac8eb2f..f4c764ee8a8 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -14,40 +14,34 @@ mod batch; mod batch_convert; +mod flow_notifier; mod pending_worker; mod region_write; +mod tables; #[cfg(test)] mod test_util; -use std::collections::{HashMap, HashSet}; use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Duration; -use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; -use api::v1::{ColumnSchema, RowInsertRequests, Rows}; -use arrow::datatypes::Schema as ArrowSchema; -use async_trait::async_trait; +use api::v1::RowInsertRequests; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; -use common_batcher::notifier::{Notifier, run_notifier}; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; -use common_runtime::spawn_global; -use common_telemetry::{error, warn}; -use datatypes::timestamp::append_timestamps; use meter_core::data::MeterRecord; use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; -use table::metadata::TableId; +use snafu::ResultExt; use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; +use crate::batcher::flow_notifier::{FlowNotifier, start_flow_notification_worker}; pub use crate::batcher::logical_table::batch::flush_batch_physical; pub use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; use crate::batcher::logical_table::pending_worker::{ @@ -57,15 +51,14 @@ pub use crate::batcher::logical_table::region_write::{ PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, }; +pub use crate::batcher::logical_table::tables::{ + PendingRowsSchemaAlterer, PendingRowsSchemaAltererRef, +}; use crate::error; use crate::error::{Error, Result}; use crate::metrics::{ FLOW_NOTIFICATION_DROPPED, PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED, PENDING_WORKERS, }; -use crate::prom_row_builder::{ - build_prom_create_table_schema_from_proto, identify_missing_columns_from_proto, - rows_to_aligned_record_batch, -}; const PHYSICAL_TABLE_KEY: &str = "physical_table"; @@ -89,34 +82,6 @@ pub fn pending_rows_batch_sync_enabled() -> bool { const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; -const MAX_CONCURRENT_FLOW_NOTIFICATIONS: NonZeroUsize = NonZeroUsize::new(8).unwrap(); - -#[async_trait] -pub trait PendingRowsSchemaAlterer: Send + Sync { - /// Batch-create multiple logical tables that are missing. - /// Each entry is `(table_name, request_schema)`. - async fn create_tables_if_missing_batch( - &self, - catalog: &str, - schema: &str, - tables: &[(&str, &[ColumnSchema])], - with_metric_engine: bool, - ctx: QueryContextRef, - ) -> Result<()>; - - /// Batch-alter multiple logical tables to add missing tag columns. - /// Each entry is `(table_name, missing_column_names)`. - async fn add_missing_prom_tag_columns_batch( - &self, - catalog: &str, - schema: &str, - tables: &[(&str, &[String])], - ctx: QueryContextRef, - ) -> Result<()>; -} - -pub type PendingRowsSchemaAltererRef = Arc; - #[derive(Debug, Clone, Hash, Eq, PartialEq)] struct BatchKey { catalog: String, @@ -125,17 +90,6 @@ struct BatchKey { skip_wal: bool, } -/// Intermediate planning state for resolving and preparing logical tables -/// before row-to-batch alignment. -struct TableResolutionPlan { - /// Resolved table schema and table id by logical table name. - region_schemas: HashMap, u32)>, - /// Missing tables that need to be created before alignment. - tables_to_create: Vec<(String, Vec)>, - /// Existing tables that need tag-column schema evolution. - tables_to_alter: Vec<(String, Vec)>, -} - // Requests can share a batch only when their write target and WAL policy match. fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { let physical_table = ctx @@ -158,7 +112,7 @@ pub struct LogicalTablePendingRowsBatcher { partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, + flow_notification_tx: FlowNotifier, flush_limiter: FlushLimiter, request_limiter: RequestLimiter, worker_channel_capacity: usize, @@ -192,8 +146,10 @@ impl LogicalTablePendingRowsBatcher { let flush_limiter = FlushLimiter::try_new(max_concurrent_flushes)?; let request_limiter = RequestLimiter::try_new(max_inflight_requests)?; - let (flow_notification_tx, flow_notification_rx) = - Notifier::try_new(flow_notification_queue_capacity.get())?; + let (flow_notification_tx, flow_notification_rx) = FlowNotifier::try_new( + flow_notification_queue_capacity.get(), + FLOW_NOTIFICATION_DROPPED.clone(), + )?; let (shutdown, _) = broadcast::channel(1); let pending_rows_batch_sync = pending_rows_batch_sync_enabled(); @@ -320,358 +276,6 @@ impl LogicalTablePendingRowsBatcher { } } -impl LogicalTablePendingRowsBatcher { - /// Converts proto `RowInsertRequests` directly into aligned `RecordBatch`es - /// in a single pass, handling table creation, schema alteration, column - /// renaming, reordering, and null-filling without building intermediate - /// RecordBatches. - async fn build_and_align_table_batches( - &self, - requests: RowInsertRequests, - ctx: &QueryContextRef, - ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { - let catalog = ctx.current_catalog().to_string(); - let schema = ctx.current_schema(); - - let (table_rows, total_rows) = Self::collect_non_empty_table_rows(requests); - if total_rows == 0 { - return Ok((Vec::new(), 0)); - } - - let unique_tables = Self::collect_unique_table_schemas(&table_rows)?; - let mut plan = self - .plan_table_resolution(&catalog, &schema, ctx, &unique_tables) - .await?; - - self.create_missing_tables_and_refresh_schemas( - &catalog, - &schema, - ctx, - &table_rows, - &mut plan, - ) - .await?; - - self.alter_tables_and_refresh_schemas(&catalog, &schema, ctx, &mut plan) - .await?; - - let aligned_batches = Self::build_aligned_batches(&table_rows, &plan.region_schemas)?; - - Ok((aligned_batches, total_rows)) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Extracts non-empty `(table_name, rows)` pairs and computes total row - /// count across the retained entries. - fn collect_non_empty_table_rows(requests: RowInsertRequests) -> (Vec<(String, Rows)>, usize) { - let mut table_rows: Vec<(String, Rows)> = Vec::with_capacity(requests.inserts.len()); - let mut total_rows = 0; - - for request in requests.inserts { - let Some(rows) = request.rows else { - continue; - }; - if rows.rows.is_empty() { - continue; - } - - total_rows += rows.rows.len(); - table_rows.push((request.table_name, rows)); - } - - (table_rows, total_rows) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Returns unique `(table_name, proto_schema)` pairs while keeping the - /// first-seen schema for duplicate table names. - fn collect_unique_table_schemas( - table_rows: &[(String, Rows)], - ) -> Result> { - let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); - let mut seen = HashSet::new(); - - for (table_name, rows) in table_rows { - if seen.insert(table_name.as_str()) { - unique_tables.push((table_name.as_str(), &rows.schema)); - } else { - // table_rows should group rows by table name. - return error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Found duplicated table name in RowInsertRequest: {}", - table_name - ), - } - .fail(); - } - } - - Ok(unique_tables) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Resolves table metadata and classifies each table into existing, - /// to-create, and to-alter groups used by subsequent DDL steps. - async fn plan_table_resolution( - &self, - catalog: &str, - schema: &str, - ctx: &QueryContextRef, - unique_tables: &[(&str, &[ColumnSchema])], - ) -> Result { - let mut plan = TableResolutionPlan { - region_schemas: HashMap::with_capacity(unique_tables.len()), - tables_to_create: Vec::new(), - tables_to_alter: Vec::new(), - }; - - let resolved_tables = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_resolve_table"]) - .start_timer(); - futures::future::join_all(unique_tables.iter().map(|(table_name, _)| { - self.catalog_manager - .table(catalog, schema, table_name, Some(ctx.as_ref())) - })) - .await - }; - - for ((table_name, rows_schema), table_result) in unique_tables.iter().zip(resolved_tables) { - let table = table_result?; - - if let Some(table) = table { - let table_info = table.table_info(); - let table_id = table_info.ident.table_id; - let region_schema = table_info.meta.schema.arrow_schema().clone(); - - let missing_columns = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_identify_missing_columns"]) - .start_timer(); - identify_missing_columns_from_proto(rows_schema, region_schema.as_ref())? - }; - if !missing_columns.is_empty() { - plan.tables_to_alter - .push(((*table_name).to_string(), missing_columns)); - } - plan.region_schemas - .insert((*table_name).to_string(), (region_schema, table_id)); - } else { - let request_schema = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_build_create_table_schema"]) - .start_timer(); - build_prom_create_table_schema_from_proto(rows_schema)? - }; - plan.tables_to_create - .push(((*table_name).to_string(), request_schema)); - } - } - - Ok(plan) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Batch-creates missing tables, refreshes their schema metadata, and - /// enqueues follow-up alters for extra tag columns discovered in later rows. - async fn create_missing_tables_and_refresh_schemas( - &self, - catalog: &str, - schema: &str, - ctx: &QueryContextRef, - table_rows: &[(String, Rows)], - plan: &mut TableResolutionPlan, - ) -> Result<()> { - if plan.tables_to_create.is_empty() { - return Ok(()); - } - - let create_refs: Vec<(&str, &[ColumnSchema])> = plan - .tables_to_create - .iter() - .map(|(name, schema)| (name.as_str(), schema.as_slice())) - .collect(); - - { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_batch_create_tables"]) - .start_timer(); - self.schema_alterer - .create_tables_if_missing_batch( - catalog, - schema, - &create_refs, - self.prom_store_with_metric_engine, - ctx.clone(), - ) - .await?; - } - - let created_table_results = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_resolve_table_after_create"]) - .start_timer(); - futures::future::join_all(plan.tables_to_create.iter().map(|(table_name, _)| { - self.catalog_manager - .table(catalog, schema, table_name, Some(ctx.as_ref())) - })) - .await - }; - - for ((table_name, _), table_result) in - plan.tables_to_create.iter().zip(created_table_results) - { - let table = table_result?.with_context(|| error::UnexpectedResultSnafu { - reason: format!( - "Table not found after pending batch create attempt: {}", - table_name - ), - })?; - let table_info = table.table_info(); - let table_id = table_info.ident.table_id; - let region_schema = table_info.meta.schema.arrow_schema().clone(); - plan.region_schemas - .insert(table_name.clone(), (region_schema, table_id)); - } - - Self::enqueue_alter_for_new_tables(table_rows, plan)?; - - Ok(()) - } -} - -impl LogicalTablePendingRowsBatcher { - /// For newly created tables, re-checks all row schemas and appends alter - /// operations when additional tag columns are still missing. - fn enqueue_alter_for_new_tables( - table_rows: &[(String, Rows)], - plan: &mut TableResolutionPlan, - ) -> Result<()> { - let created_tables: HashSet<&str> = plan - .tables_to_create - .iter() - .map(|(table_name, _)| table_name.as_str()) - .collect(); - - for (table_name, rows) in table_rows { - if !created_tables.contains(table_name.as_str()) { - continue; - } - - let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { - continue; - }; - - let missing_columns = identify_missing_columns_from_proto(&rows.schema, region_schema)?; - if missing_columns.is_empty() - || plan - .tables_to_alter - .iter() - .any(|(existing_name, _)| existing_name == table_name) - { - continue; - } - - plan.tables_to_alter - .push((table_name.clone(), missing_columns)); - } - - Ok(()) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Batch-alters tables that have missing tag columns and refreshes the - /// in-memory schema map used for row alignment. - async fn alter_tables_and_refresh_schemas( - &self, - catalog: &str, - schema: &str, - ctx: &QueryContextRef, - plan: &mut TableResolutionPlan, - ) -> Result<()> { - if plan.tables_to_alter.is_empty() { - return Ok(()); - } - - let alter_refs: Vec<(&str, &[String])> = plan - .tables_to_alter - .iter() - .map(|(name, cols)| (name.as_str(), cols.as_slice())) - .collect(); - { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_batch_add_missing_columns"]) - .start_timer(); - self.schema_alterer - .add_missing_prom_tag_columns_batch(catalog, schema, &alter_refs, ctx.clone()) - .await?; - } - - let altered_table_results = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_resolve_table_after_schema_alter"]) - .start_timer(); - futures::future::join_all(plan.tables_to_alter.iter().map(|(table_name, _)| { - self.catalog_manager - .table(catalog, schema, table_name, Some(ctx.as_ref())) - })) - .await - }; - - for ((table_name, _), table_result) in - plan.tables_to_alter.iter().zip(altered_table_results) - { - let table = table_result?.with_context(|| error::UnexpectedResultSnafu { - reason: format!( - "Table not found after pending batch schema alter: {}", - table_name - ), - })?; - let table_info = table.table_info(); - let table_id = table_info.ident.table_id; - let refreshed_region_schema = table_info.meta.schema.arrow_schema().clone(); - plan.region_schemas - .insert(table_name.clone(), (refreshed_region_schema, table_id)); - } - - Ok(()) - } -} - -impl LogicalTablePendingRowsBatcher { - /// Converts proto rows to `RecordBatch` values aligned to resolved region - /// schemas and returns `(table_name, table_id, batch)` tuples. - fn build_aligned_batches( - table_rows: &[(String, Rows)], - region_schemas: &HashMap, u32)>, - ) -> Result> { - let mut aligned_batches = Vec::with_capacity(table_rows.len()); - for (table_name, rows) in table_rows { - let (region_schema, table_id) = - region_schemas.get(table_name).cloned().with_context(|| { - error::UnexpectedResultSnafu { - reason: format!("Region schema not resolved for table: {}", table_name), - } - })?; - - let record_batch = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["align_rows_to_record_batch"]) - .start_timer(); - rows_to_aligned_record_batch(rows, region_schema.as_ref())? - }; - aligned_batches.push((table_name.clone(), table_id, record_batch)); - } - - Ok(aligned_batches) - } -} - impl LogicalTablePendingRowsBatcher { async fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker { let (tx, receiver) = self @@ -721,206 +325,18 @@ impl Drop for LogicalTablePendingRowsBatcher { } } -fn extract_timestamps(table_batch: &TableBatch) -> Vec { - let mut timestamps = Vec::with_capacity(table_batch.row_count); - for batch in &table_batch.batches { - let timestamp_column = batch.batch.column(batch.timestamp_index); - let Some(()) = append_timestamps(timestamp_column, &mut timestamps) else { - error!( - "Failed to extract timestamps from record batch, table_id: {}, timestamp_index: {}", - table_batch.table_id, batch.timestamp_index - ); - continue; - }; - } - timestamps -} - -struct FlowNotification { - table_id: TableId, - timestamps: Vec, -} - -fn try_enqueue_flow_notification( - tx: &Notifier, - notification: FlowNotification, -) -> bool { - match tx.try_notify(notification) { - Ok(()) => true, - Err(mpsc::error::TrySendError::Full(notification)) => { - FLOW_NOTIFICATION_DROPPED.with_label_values(&["full"]).inc(); - warn!( - "Dropping flow notification because queue is full, table_id: {}, queue_capacity: {}", - notification.table_id, - tx.max_capacity() - ); - false - } - Err(mpsc::error::TrySendError::Closed(notification)) => { - FLOW_NOTIFICATION_DROPPED - .with_label_values(&["closed"]) - .inc(); - error!( - "Dropping flow notification because queue is closed, table_id: {}, queue_capacity: {}", - notification.table_id, - tx.max_capacity() - ); - false - } - } -} - -fn enqueue_flow_notifications(table_batches: Vec, tx: &Notifier) { - for table_batch in table_batches { - let timestamps = extract_timestamps(&table_batch); - if timestamps.is_empty() { - continue; - } - try_enqueue_flow_notification( - tx, - FlowNotification { - table_id: table_batch.table_id, - timestamps, - }, - ); - } -} - -async fn handle_flow_notification( - notification: FlowNotification, - table_flownode_set_cache: TableFlownodeSetCacheRef, - node_manager: NodeManagerRef, -) { - let table_id = notification.table_id; - let flownodes = match table_flownode_set_cache.get(table_id).await { - Ok(Some(flownodes)) => flownodes, - Ok(None) => return, - Err(e) => { - error!(e; "Failed to get flownodes for table id: {}", table_id); - return; - } - }; - let peers = flownodes.values().cloned().collect::>(); - - for peer in peers { - if let Err(e) = node_manager - .flownode(&peer) - .await - .handle_mark_window_dirty(DirtyWindowRequests { - requests: vec![DirtyWindowRequest { - table_id, - timestamps: notification.timestamps.clone(), - time_ranges: Vec::new(), - }], - }) - .await - { - error!( - e; - "Failed to mark timestamps as dirty, table_id: {}, peer_id: {}, peer_addr: {}", - table_id, - peer.id, - peer.addr - ); - } - } -} - -fn start_flow_notification_worker( - notification_rx: mpsc::Receiver, - table_flownode_set_cache: TableFlownodeSetCacheRef, - node_manager: NodeManagerRef, -) { - spawn_global(async move { - run_notifier( - notification_rx, - MAX_CONCURRENT_FLOW_NOTIFICATIONS, - |notification| { - let table_flownode_set_cache = table_flownode_set_cache.clone(); - let node_manager = node_manager.clone(); - handle_flow_notification(notification, table_flownode_set_cache, node_manager) - }, - ) - .await; - }); -} - -#[cfg(test)] -fn notify_flow_dirty_windows_after_flush( - table_batches: Vec, - table_flownode_set_cache: TableFlownodeSetCacheRef, - node_manager: NodeManagerRef, -) { - let (tx, rx) = Notifier::try_new(table_batches.len().max(1)).unwrap(); - start_flow_notification_worker(rx, table_flownode_set_cache, node_manager); - enqueue_flow_notifications(table_batches, &tx); -} - #[cfg(test)] mod tests { - use std::any::Any; - use std::sync::{Arc, Mutex}; - use std::time::Duration; + use std::sync::Arc; - use api::v1::meta::Peer; - use api::v1::value::ValueData; - use api::v1::{ - ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, - Value, - }; use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; - use async_trait::async_trait; - use common_batcher::notifier::Notifier; - use common_meta::cache::new_table_flownode_set_cache; - use common_meta::error::Result as MetaResult; - use common_meta::instruction::{CacheIdent, CreateFlow}; - use common_meta::kv_backend::{KvBackend, TxnService}; - use common_meta::node_manager::NodeManagerRef; - use common_meta::rpc::store::{ - BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, - BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, - PutResponse, RangeRequest, RangeResponse, - }; use common_query::prelude::greptime_timestamp; - use moka::future::CacheBuilder; - use tokio::sync::{Notify, mpsc, oneshot}; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; - use crate::batcher::logical_table::test_util::{ - FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch, - mock_table_flownode_cache, mock_tag_batch, - }; - use crate::batcher::logical_table::{ - LogicalTablePendingRowsBatcher, extract_timestamps, notify_flow_dirty_windows_after_flush, - try_enqueue_flow_notification, - }; - 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 { - Rows { - schema: vec![ColumnSchema { - column_name: schema_name.to_string(), - ..Default::default() - }], - rows: (0..row_count).map(|_| Row { values: vec![] }).collect(), - } - } - - fn mock_timestamp_batch(timestamps: Vec>) -> RecordBatchWithTsIdx { - let batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![Field::new( - greptime_timestamp(), - ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - true, - )])), - vec![Arc::new(TimestampMillisecondArray::from(timestamps))], - ) - .unwrap(); - RecordBatchWithTsIdx::try_new(batch, 0).unwrap() - } + use crate::batcher::logical_table::flow_notifier::extract_timestamps; + use crate::batcher::logical_table::test_util::{mock_aligned_tag_batch, mock_tag_batch}; #[test] fn test_extract_timestamps_appends_non_null_batches_in_order() { @@ -1014,327 +430,18 @@ mod tests { ); } - #[test] - fn test_extract_timestamps_uses_aligned_custom_timestamp_index() { - let rows = Rows { - schema: vec![ - ColumnSchema { - column_name: greptime_timestamp().to_string(), - datatype: ColumnDataType::TimestampMillisecond as i32, - semantic_type: SemanticType::Timestamp as i32, - ..Default::default() - }, - ColumnSchema { - column_name: "host".to_string(), - datatype: ColumnDataType::String as i32, - semantic_type: SemanticType::Tag as i32, - ..Default::default() - }, - ColumnSchema { - column_name: "greptime_value".to_string(), - datatype: ColumnDataType::Float64 as i32, - semantic_type: SemanticType::Field as i32, - ..Default::default() - }, - ], - rows: vec![ - Row { - values: vec![ - Value { - value_data: Some(ValueData::TimestampMillisecondValue(1000)), - }, - Value { - value_data: Some(ValueData::StringValue("host-1".to_string())), - }, - Value { - value_data: Some(ValueData::F64Value(1.0)), - }, - ], - }, - Row { - values: vec![ - Value { - value_data: Some(ValueData::TimestampMillisecondValue(2000)), - }, - Value { - value_data: Some(ValueData::StringValue("host-2".to_string())), - }, - Value { - value_data: Some(ValueData::F64Value(2.0)), - }, - ], - }, - ], - }; - let target_schema = ArrowSchema::new(vec![ - Field::new("host", ArrowDataType::Utf8, true), - Field::new( - "timestamp", + fn mock_timestamp_batch(timestamps: Vec>) -> RecordBatchWithTsIdx { + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + greptime_timestamp(), ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), - Field::new("greptime_value", ArrowDataType::Float64, true), - ]); - let batch = rows_to_aligned_record_batch(&rows, &target_schema).unwrap(); - assert_eq!(1, batch.timestamp_index); - let table_batch = TableBatch { - table_name: "cpu".to_string(), - table_id: 42, - row_count: batch.batch.num_rows(), - batches: vec![batch], - }; - - assert_eq!(vec![1000, 2000], extract_timestamps(&table_batch)); + true, + )])), + vec![Arc::new(TimestampMillisecondArray::from(timestamps))], + ) + .unwrap(); + RecordBatchWithTsIdx::try_new(batch, 0).unwrap() } - - #[test] - fn test_flow_notification_queue_drops_when_full() { - let (tx, mut rx) = Notifier::try_new(1).unwrap(); - let notification = |table_id| crate::batcher::logical_table::FlowNotification { - table_id, - timestamps: vec![table_id as i64], - }; - let dropped = FLOW_NOTIFICATION_DROPPED.with_label_values(&["full"]); - let dropped_before = dropped.get(); - - assert!(try_enqueue_flow_notification(&tx, notification(1))); - assert!(!try_enqueue_flow_notification(&tx, notification(2))); - - assert_eq!(1, rx.try_recv().unwrap().table_id); - assert_eq!(dropped_before + 1, dropped.get()); - - let closed = FLOW_NOTIFICATION_DROPPED.with_label_values(&["closed"]); - let closed_before = closed.get(); - drop(rx); - assert!(!try_enqueue_flow_notification(&tx, notification(3))); - assert_eq!(closed_before + 1, closed.get()); - } - - #[test] - fn test_collect_non_empty_table_rows_filters_empty_payloads() { - let requests = RowInsertRequests { - inserts: vec![ - RowInsertRequest { - table_name: "cpu".to_string(), - rows: Some(mock_rows(2, "host")), - }, - RowInsertRequest { - table_name: "mem".to_string(), - rows: Some(mock_rows(0, "host")), - }, - RowInsertRequest { - table_name: "disk".to_string(), - rows: None, - }, - ], - }; - - let (table_rows, total_rows) = - LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); - - assert_eq!(2, total_rows); - assert_eq!(1, table_rows.len()); - assert_eq!("cpu", table_rows[0].0); - assert_eq!(2, table_rows[0].1.rows.len()); - } - - struct BlockingRangeKvBackend { - range_started: Mutex>>, - range_release: Arc, - } - - impl TxnService for BlockingRangeKvBackend { - type Error = common_meta::error::Error; - } - - #[async_trait] - impl KvBackend for BlockingRangeKvBackend { - fn name(&self) -> &str { - "blocking_range" - } - - fn as_any(&self) -> &dyn Any { - self - } - - async fn range(&self, _req: RangeRequest) -> MetaResult { - let range_started = self.range_started.lock().unwrap().take(); - if let Some(range_started) = range_started { - let _ = range_started.send(()); - self.range_release.notified().await; - } - Ok(RangeResponse { - kvs: Vec::new(), - more: false, - }) - } - - async fn put(&self, _req: PutRequest) -> MetaResult { - unimplemented!() - } - - async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult { - unimplemented!() - } - - async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult { - unimplemented!() - } - - async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult { - unimplemented!() - } - - async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult { - unimplemented!() - } - } - - #[tokio::test] - async fn test_flow_notifications_do_not_block_on_previous_table_cache_lookup() { - let blocked_table_id = 41; - let cached_table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let (range_started_tx, range_started_rx) = oneshot::channel(); - let range_release = Arc::new(Notify::new()); - let cache = Arc::new(new_table_flownode_set_cache( - "test".to_string(), - CacheBuilder::new(2).build(), - Arc::new(BlockingRangeKvBackend { - range_started: Mutex::new(Some(range_started_tx)), - range_release: range_release.clone(), - }), - )); - cache - .invalidate(&[CacheIdent::CreateFlow(CreateFlow { - flow_id: 1, - source_table_ids: vec![cached_table_id], - partition_to_peer_mapping: vec![(0, peer)], - })]) - .await - .unwrap(); - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let table_batches = vec![ - TableBatch { - table_name: "blocked".to_string(), - table_id: blocked_table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }, - TableBatch { - table_name: "cached".to_string(), - table_id: cached_table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], - row_count: 1, - }, - ]; - - notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); - - tokio::time::timeout(Duration::from_secs(1), range_started_rx) - .await - .unwrap() - .unwrap(); - let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) - .await - .unwrap() - .unwrap(); - assert_eq!(cached_table_id, requests.requests[0].table_id); - range_release.notify_one(); - } - - #[tokio::test] - async fn test_successful_flush_notifies_flownode_with_logical_table_timestamps() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - - notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); - - let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) - .await - .unwrap() - .unwrap(); - assert_eq!( - vec![api::v1::flow::DirtyWindowRequest { - table_id, - timestamps: vec![1000], - time_ranges: Vec::new(), - }], - requests.requests - ); - assert!( - tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) - .await - .is_err() - ); - } - - #[tokio::test] - async fn test_successful_flush_coalesces_logical_batches_per_flownode() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![ - mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0), - mock_aligned_tag_batch("tag1", "host-1", 2000, 2.0), - ], - row_count: 2, - }]; - - notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); - - let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) - .await - .unwrap() - .unwrap(); - assert_eq!( - vec![api::v1::flow::DirtyWindowRequest { - table_id, - timestamps: vec![1000, 2000], - time_ranges: Vec::new(), - }], - requests.requests - ); - assert!( - tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) - .await - .is_err() - ); - } - #[test] fn test_batch_key_groups_by_skip_wal() { use std::collections::HashMap; diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs index af9094ef724..ac2db137de1 100644 --- a/src/servers/src/batcher/logical_table/batch.rs +++ b/src/servers/src/batcher/logical_table/batch.rs @@ -18,7 +18,6 @@ use std::time::Instant; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; -use common_batcher::notifier::Notifier; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; use common_telemetry::{debug, warn}; @@ -26,9 +25,12 @@ 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::pending_worker::FlushWaiter; use crate::batcher::logical_table::region_write::{ CatalogManagerPhysicalFlushAdapter, NodeManagerPhysicalFlushAdapter, @@ -36,9 +38,6 @@ use crate::batcher::logical_table::region_write::{ PhysicalFlushPartitionProvider, encode_region_write_requests, flush_region_writes_concurrently, plan_region_batches, resolve_region_targets, }; -use crate::batcher::logical_table::{ - FlowNotification, PHYSICAL_TABLE_KEY, enqueue_flow_notifications, -}; use crate::error; use crate::error::Result; use crate::metrics::{ @@ -59,7 +58,7 @@ pub(in crate::batcher::logical_table) async fn spawn_flush( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, + flow_notification_tx: FlowNotifier, flush_limiter: FlushLimiter, ) { match flush_limiter.acquire().await { @@ -95,7 +94,7 @@ pub(in crate::batcher::logical_table) async fn flush_batch_with_managers( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, + flow_notification_tx: FlowNotifier, ) { let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; let node_requester = NodeManagerPhysicalFlushAdapter { @@ -117,7 +116,7 @@ pub(in crate::batcher::logical_table) async fn flush_batch( partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), - flow_notification_tx: Notifier, + flow_notification_tx: FlowNotifier, ) { let Batch { table_batches, @@ -283,7 +282,6 @@ mod tests { use arrow::record_batch::RecordBatch; use async_trait::async_trait; use catalog::error::Result as CatalogResult; - use common_batcher::notifier::Notifier; use common_batcher::request_limiter::RequestLimiter; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; @@ -295,6 +293,7 @@ mod tests { use table::test_util::table_info::test_table_info; use tokio::sync::{mpsc, oneshot}; + use crate::batcher::flow_notifier::{FlowNotifier, start_flow_notification_worker}; use crate::batcher::logical_table::batch::{ Batch, flush_batch, flush_batch_physical, notify_waiters, }; @@ -304,13 +303,286 @@ mod tests { PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, }; - use crate::batcher::logical_table::start_flow_notification_worker; use crate::batcher::logical_table::test_util::{ FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch, - mock_table_flownode_cache, }; + use crate::batcher::test_util::mock_table_flownode_cache; use crate::error; use crate::error::Error; + use crate::metrics::FLOW_NOTIFICATION_DROPPED; + + #[tokio::test] + async fn test_flush_batch_notifies_flownode_after_successful_physical_write() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); + let ctx = session::context::QueryContext::arc(); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let writes = Arc::new(AtomicUsize::new(0)); + + flush_batch( + Batch { + table_batches, + total_row_count: 1, + db_string: ctx.get_db_string(), + ctx, + waiters: Vec::new(), + }, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &MockFlushNodeRequester { + writes: writes.clone(), + fail: false, + }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + flow_notification_tx, + ) + .await; + + assert_eq!(1, writes.load(Ordering::SeqCst)); + let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + vec![api::v1::flow::DirtyWindowRequest { + table_id, + timestamps: vec![1000], + time_ranges: Vec::new(), + }], + requests.requests + ); + } + + #[tokio::test] + async fn test_flush_batch_does_not_notify_flownode_after_physical_write_error() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); + let ctx = session::context::QueryContext::arc(); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let writes = Arc::new(AtomicUsize::new(0)); + + flush_batch( + Batch { + table_batches, + total_row_count: 1, + db_string: ctx.get_db_string(), + ctx, + waiters: Vec::new(), + }, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &MockFlushNodeRequester { + writes: writes.clone(), + fail: true, + }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + flow_notification_tx, + ) + .await; + + assert_eq!(1, writes.load(Ordering::SeqCst)); + assert!( + tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_cancelled_waiter_retains_request_slot_until_notification() { + let limiter = RequestLimiter::try_new(1).unwrap(); + let (response_tx, response_rx) = oneshot::channel(); + let waiter = FlushWaiter { + response_tx, + _permit: limiter.acquire().await.unwrap(), + }; + drop(response_rx); + let next = limiter.acquire(); + tokio::pin!(next); + assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await); + notify_waiters(vec![waiter], Ok(())); + let _permit = next.await.unwrap(); + } + + #[tokio::test] + async fn test_flush_batch_physical_uses_mockable_trait_dependencies() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .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)); + } + + #[tokio::test] + async fn test_flush_batch_physical_returns_actual_affected_rows() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let ctx = session::context::QueryContext::arc(); + + let affected_rows = flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + region_leader_calls: Arc::new(AtomicUsize::new(0)), + }, + &AffectedRowsFlushNodeRequester { affected_rows: 7 }, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .unwrap(); + + assert_eq!(7, affected_rows); + } + + #[tokio::test] + async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() { + let table_batches = vec![TableBatch { + table_name: "t1".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + let err = flush_batch_physical( + &table_batches, + "missing_phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { table: None }, + ) + .await + .unwrap_err(); + + assert!( + err.to_string() + .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)); + } + + #[tokio::test] + async fn test_flush_batch_physical_aborts_immediately_on_transform_error() { + let table_batches = vec![ + TableBatch { + table_name: "broken".to_string(), + table_id: 11, + batches: vec![mock_aligned_tag_batch("unknown_tag", "host-1", 1000, 1.0)], + row_count: 1, + }, + TableBatch { + table_name: "healthy".to_string(), + table_id: 12, + batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], + row_count: 1, + }, + ]; + let partition_calls = Arc::new(AtomicUsize::new(0)); + let leader_calls = Arc::new(AtomicUsize::new(0)); + let node = MockFlushNodeRequester::default(); + let ctx = session::context::QueryContext::arc(); + + let err = flush_batch_physical( + &table_batches, + "phy", + &ctx, + &MockFlushPartitionProvider { + partition_rule_calls: partition_calls.clone(), + region_leader_calls: leader_calls.clone(), + }, + &node, + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + ) + .await + .unwrap_err(); + + 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)); + } fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata { let schema = Arc::new( @@ -447,177 +719,12 @@ mod tests { fn mock_flow_notification_sender( cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, - ) -> Notifier { - let (tx, rx) = Notifier::try_new(16).unwrap(); + ) -> FlowNotifier { + let (tx, rx) = FlowNotifier::try_new(16, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); start_flow_notification_worker(rx, cache, node_manager); tx } - #[tokio::test] - async fn test_flush_batch_notifies_flownode_after_successful_physical_write() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); - let ctx = session::context::QueryContext::arc(); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let writes = Arc::new(AtomicUsize::new(0)); - - flush_batch( - Batch { - table_batches, - total_row_count: 1, - db_string: ctx.get_db_string(), - ctx, - waiters: Vec::new(), - }, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &MockFlushNodeRequester { - writes: writes.clone(), - fail: false, - }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - flow_notification_tx, - ) - .await; - - assert_eq!(1, writes.load(Ordering::SeqCst)); - let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) - .await - .unwrap() - .unwrap(); - assert_eq!( - vec![api::v1::flow::DirtyWindowRequest { - table_id, - timestamps: vec![1000], - time_ranges: Vec::new(), - }], - requests.requests - ); - } - - #[tokio::test] - async fn test_flush_batch_does_not_notify_flownode_after_physical_write_error() { - let table_id = 42; - let peer = Peer { - id: 7, - addr: "flow-7".to_string(), - }; - let cache = mock_table_flownode_cache(table_id, peer).await; - let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); - let _requests_tx = requests_tx.clone(); - let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { - flownode: Arc::new(RecordingFlownode { requests_tx }), - }); - let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone()); - let ctx = session::context::QueryContext::arc(); - let table_batches = vec![TableBatch { - table_name: "cpu".to_string(), - table_id, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let writes = Arc::new(AtomicUsize::new(0)); - - flush_batch( - Batch { - table_batches, - total_row_count: 1, - db_string: ctx.get_db_string(), - ctx, - waiters: Vec::new(), - }, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &MockFlushNodeRequester { - writes: writes.clone(), - fail: true, - }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - flow_notification_tx, - ) - .await; - - assert_eq!(1, writes.load(Ordering::SeqCst)); - assert!( - tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) - .await - .is_err() - ); - } - - #[tokio::test] - async fn test_cancelled_waiter_retains_request_slot_until_notification() { - let limiter = RequestLimiter::try_new(1).unwrap(); - let (response_tx, response_rx) = oneshot::channel(); - let waiter = FlushWaiter { - response_tx, - _permit: limiter.acquire().await.unwrap(), - }; - drop(response_rx); - let next = limiter.acquire(); - tokio::pin!(next); - assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await); - notify_waiters(vec![waiter], Ok(())); - let _permit = next.await.unwrap(); - } - - #[tokio::test] - async fn test_flush_batch_physical_uses_mockable_trait_dependencies() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .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)); - } - #[derive(Default)] struct AffectedRowsFlushNodeRequester { affected_rows: usize, @@ -633,112 +740,4 @@ mod tests { Ok(RegionResponse::new(self.affected_rows)) } } - - #[tokio::test] - async fn test_flush_batch_physical_returns_actual_affected_rows() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let ctx = session::context::QueryContext::arc(); - - let affected_rows = flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: Arc::new(AtomicUsize::new(0)), - region_leader_calls: Arc::new(AtomicUsize::new(0)), - }, - &AffectedRowsFlushNodeRequester { affected_rows: 7 }, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .unwrap(); - - assert_eq!(7, affected_rows); - } - - #[tokio::test] - async fn test_flush_batch_physical_stops_before_partition_and_node_when_table_missing() { - let table_batches = vec![TableBatch { - table_name: "t1".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], - row_count: 1, - }]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - let err = flush_batch_physical( - &table_batches, - "missing_phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { table: None }, - ) - .await - .unwrap_err(); - - assert!( - err.to_string() - .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)); - } - - #[tokio::test] - async fn test_flush_batch_physical_aborts_immediately_on_transform_error() { - let table_batches = vec![ - TableBatch { - table_name: "broken".to_string(), - table_id: 11, - batches: vec![mock_aligned_tag_batch("unknown_tag", "host-1", 1000, 1.0)], - row_count: 1, - }, - TableBatch { - table_name: "healthy".to_string(), - table_id: 12, - batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], - row_count: 1, - }, - ]; - let partition_calls = Arc::new(AtomicUsize::new(0)); - let leader_calls = Arc::new(AtomicUsize::new(0)); - let node = MockFlushNodeRequester::default(); - let ctx = session::context::QueryContext::arc(); - - let err = flush_batch_physical( - &table_batches, - "phy", - &ctx, - &MockFlushPartitionProvider { - partition_rule_calls: partition_calls.clone(), - region_leader_calls: leader_calls.clone(), - }, - &node, - &MockFlushCatalogProvider { - table: Some(mock_physical_table_metadata(1024)), - }, - ) - .await - .unwrap_err(); - - 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)); - } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index eba2335df03..28464eb2d3b 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -290,9 +290,10 @@ mod tests { use std::collections::{HashMap, HashSet}; use std::sync::Arc; - use arrow::array::{BinaryArray, StringArray, TimestampMillisecondArray}; + use arrow::array::{Array, BinaryArray, 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::{ @@ -501,7 +502,6 @@ mod tests { #[test] fn test_modify_batch_sparse_with_taxonomy_per_batch() { use arrow::array::BinaryArray; - use metric_engine::batch_modifier::modify_batch_sparse; let schema1 = Arc::new(ArrowSchema::new(vec![ Field::new( diff --git a/src/servers/src/batcher/logical_table/flow_notifier.rs b/src/servers/src/batcher/logical_table/flow_notifier.rs new file mode 100644 index 00000000000..83ac6dae728 --- /dev/null +++ b/src/servers/src/batcher/logical_table/flow_notifier.rs @@ -0,0 +1,326 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +use common_meta::cache::TableFlownodeSetCacheRef; +#[cfg(test)] +use common_meta::node_manager::NodeManagerRef; +use common_telemetry::error; +use datatypes::timestamp::append_timestamps; + +use crate::batcher::flow_notifier::FlowNotifier; +#[cfg(test)] +use crate::batcher::flow_notifier::start_flow_notification_worker; +use crate::batcher::flow_sender::FlowNotification; +use crate::batcher::logical_table::batch_convert::TableBatch; +#[cfg(test)] +use crate::metrics::FLOW_NOTIFICATION_DROPPED; + +pub(in crate::batcher::logical_table) fn extract_timestamps(table_batch: &TableBatch) -> Vec { + let mut timestamps = Vec::with_capacity(table_batch.row_count); + for batch in &table_batch.batches { + let timestamp_column = batch.batch.column(batch.timestamp_index); + let Some(()) = append_timestamps(timestamp_column, &mut timestamps) else { + error!( + "Failed to extract timestamps from record batch, table_id: {}, timestamp_index: {}", + table_batch.table_id, batch.timestamp_index + ); + continue; + }; + } + timestamps +} + +pub(in crate::batcher::logical_table) fn enqueue_flow_notifications( + table_batches: Vec, + tx: &FlowNotifier, +) { + for table_batch in table_batches { + let timestamps = extract_timestamps(&table_batch); + if timestamps.is_empty() { + continue; + } + tx.try_notify(FlowNotification { + table_id: table_batch.table_id, + timestamps, + }); + } +} + +#[cfg(test)] +pub(in crate::batcher::logical_table) fn notify_flow_dirty_windows_after_flush( + table_batches: Vec, + table_flownode_set_cache: TableFlownodeSetCacheRef, + node_manager: NodeManagerRef, +) { + let (tx, rx) = FlowNotifier::try_new( + table_batches.len().max(1), + FLOW_NOTIFICATION_DROPPED.clone(), + ) + .unwrap(); + start_flow_notification_worker(rx, table_flownode_set_cache, node_manager); + enqueue_flow_notifications(table_batches, &tx); +} + +#[cfg(test)] +mod tests { + use std::any::Any; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use api::v1::meta::Peer; + use async_trait::async_trait; + use common_meta::cache::new_table_flownode_set_cache; + use common_meta::error::Result as MetaResult; + use common_meta::instruction::{CacheIdent, CreateFlow}; + use common_meta::kv_backend::{KvBackend, TxnService}; + use common_meta::node_manager::NodeManagerRef; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, + PutResponse, RangeRequest, RangeResponse, + }; + use moka::future::CacheBuilder; + use tokio::sync::{Notify, mpsc, oneshot}; + + use crate::batcher::flow_notifier::FlowNotifier; + use crate::batcher::flow_sender::FlowNotification; + use crate::batcher::logical_table::batch_convert::TableBatch; + use crate::batcher::logical_table::flow_notifier::notify_flow_dirty_windows_after_flush; + use crate::batcher::logical_table::test_util::{ + FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch, + }; + use crate::batcher::test_util::mock_table_flownode_cache; + use crate::metrics::FLOW_NOTIFICATION_DROPPED; + + #[test] + fn test_flow_notification_queue_drops_when_full() { + let (tx, mut rx) = FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); + let notification = |table_id| FlowNotification { + table_id, + timestamps: vec![table_id as i64], + }; + let dropped = FLOW_NOTIFICATION_DROPPED.with_label_values(&["full"]); + let dropped_before = dropped.get(); + + assert!(tx.try_notify(notification(1))); + assert!(!tx.try_notify(notification(2))); + + assert_eq!(1, rx.try_recv().unwrap().table_id); + assert_eq!(dropped_before + 1, dropped.get()); + + let closed = FLOW_NOTIFICATION_DROPPED.with_label_values(&["closed"]); + let closed_before = closed.get(); + drop(rx); + assert!(!tx.try_notify(notification(3))); + assert_eq!(closed_before + 1, closed.get()); + } + + #[tokio::test] + async fn test_flow_notifications_do_not_block_on_previous_table_cache_lookup() { + let blocked_table_id = 41; + let cached_table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let (range_started_tx, range_started_rx) = oneshot::channel(); + let range_release = Arc::new(Notify::new()); + let cache = Arc::new(new_table_flownode_set_cache( + "test".to_string(), + CacheBuilder::new(2).build(), + Arc::new(BlockingRangeKvBackend { + range_started: Mutex::new(Some(range_started_tx)), + range_release: range_release.clone(), + }), + )); + cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![cached_table_id], + partition_to_peer_mapping: vec![(0, peer)], + })]) + .await + .unwrap(); + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let table_batches = vec![ + TableBatch { + table_name: "blocked".to_string(), + table_id: blocked_table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }, + TableBatch { + table_name: "cached".to_string(), + table_id: cached_table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)], + row_count: 1, + }, + ]; + + notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); + + tokio::time::timeout(Duration::from_secs(1), range_started_rx) + .await + .unwrap() + .unwrap(); + let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(cached_table_id, requests.requests[0].table_id); + range_release.notify_one(); + } + + #[tokio::test] + async fn test_successful_flush_notifies_flownode_with_logical_table_timestamps() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)], + row_count: 1, + }]; + + notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); + + let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + vec![api::v1::flow::DirtyWindowRequest { + table_id, + timestamps: vec![1000], + time_ranges: Vec::new(), + }], + requests.requests + ); + assert!( + tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_successful_flush_coalesces_logical_batches_per_flownode() { + let table_id = 42; + let peer = Peer { + id: 7, + addr: "flow-7".to_string(), + }; + let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await; + let (requests_tx, mut requests_rx) = mpsc::unbounded_channel(); + let _requests_tx = requests_tx.clone(); + let node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager { + flownode: Arc::new(RecordingFlownode { requests_tx }), + }); + let table_batches = vec![TableBatch { + table_name: "cpu".to_string(), + table_id, + batches: vec![ + mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0), + mock_aligned_tag_batch("tag1", "host-1", 2000, 2.0), + ], + row_count: 2, + }]; + + notify_flow_dirty_windows_after_flush(table_batches, cache, node_manager); + + let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + vec![api::v1::flow::DirtyWindowRequest { + table_id, + timestamps: vec![1000, 2000], + time_ranges: Vec::new(), + }], + requests.requests + ); + assert!( + tokio::time::timeout(Duration::from_millis(50), requests_rx.recv()) + .await + .is_err() + ); + } + + struct BlockingRangeKvBackend { + range_started: Mutex>>, + range_release: Arc, + } + + impl TxnService for BlockingRangeKvBackend { + type Error = common_meta::error::Error; + } + + #[async_trait] + impl KvBackend for BlockingRangeKvBackend { + fn name(&self) -> &str { + "blocking_range" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn range(&self, _req: RangeRequest) -> MetaResult { + let range_started = self.range_started.lock().unwrap().take(); + if let Some(range_started) = range_started { + let _ = range_started.send(()); + self.range_release.notified().await; + } + Ok(RangeResponse { + kvs: Vec::new(), + more: false, + }) + } + + async fn put(&self, _req: PutRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult { + unimplemented!() + } + + async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult { + unimplemented!() + } + } +} diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs index e7bd33c1975..4a83d8a87f4 100644 --- a/src/servers/src/batcher/logical_table/pending_worker.rs +++ b/src/servers/src/batcher/logical_table/pending_worker.rs @@ -20,7 +20,6 @@ use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::FlushTrigger; use common_batcher::flush_policy::timing::TimingFlushPolicy; -use common_batcher::notifier::Notifier; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::worker_registry::WorkerRegistry; use common_meta::node_manager::NodeManagerRef; @@ -30,9 +29,10 @@ use session::context::QueryContextRef; use table::metadata::TableId; 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_convert::{RecordBatchWithTsIdx, TableBatch}; -use crate::batcher::logical_table::{BatchKey, FlowNotification}; use crate::error::Error; use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS}; @@ -102,7 +102,7 @@ pub(in crate::batcher::logical_table) fn start_worker( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: Notifier, + flow_notification_tx: FlowNotifier, worker_idle_timeout: Duration, flush_policy: TimingFlushPolicy, flush_limiter: FlushLimiter, @@ -280,7 +280,6 @@ mod tests { use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::FlushTrigger; use common_batcher::flush_policy::timing::TimingFlushPolicy; - use common_batcher::notifier::Notifier; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::new_table_route_cache; @@ -292,6 +291,7 @@ mod tests { use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; use tokio::time::advance; + use crate::batcher::flow_notifier::FlowNotifier; use crate::batcher::logical_table::BatchKey; use crate::batcher::logical_table::batch_convert::TableBatch; use crate::batcher::logical_table::pending_worker::{ @@ -302,6 +302,7 @@ mod tests { ConcurrentMockNodeManager, mock_aligned_tag_batch, }; use crate::error::Error; + use crate::metrics::FLOW_NOTIFICATION_DROPPED; #[tokio::test] async fn test_drain_batch_takes_initialized_pending_batch_from_option() { @@ -466,66 +467,6 @@ mod tests { assert!(!should_close_worker_on_idle_timeout(0, 1)); } - const WORKER_TEST_TIMEOUT: Duration = Duration::from_secs(30); - - async fn submit_mock_worker_batch( - worker_tx: &mpsc::Sender, - total_rows: usize, - timestamp: i64, - ) -> oneshot::Receiver>> { - let (response_tx, response_rx) = oneshot::channel(); - let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); - worker_tx - .send(WorkerCommand::Submit { - table_batches: vec![( - "cpu".to_string(), - 42, - mock_aligned_tag_batch("tag1", "host-1", timestamp, 1.0), - )], - total_rows, - ctx: session::context::QueryContext::arc(), - response_tx, - _permit: Arc::new(permit), - }) - .await - .unwrap(); - - // The channel is FIFO, so the ack proves the worker has dequeued and - // processed the submission (anchoring the flush deadline) before the - // caller advances virtual time. - let (ack_tx, ack_rx) = oneshot::channel(); - worker_tx.send(WorkerCommand::Ack { ack_tx }).await.unwrap(); - ack_rx - .await - .expect("worker exited before acking the submitted batch"); - - response_rx - } - - async fn receive_mock_flush_result( - response_rx: oneshot::Receiver>>, - context: &str, - ) -> std::result::Result<(), Arc> { - // Under paused time the timeout auto-advances the clock and fires - // deterministically if the flush never completes. - tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx) - .await - .unwrap_or_else(|_| panic!("{context}")) - .expect("flush result channel closed without a result") - } - - fn assert_missing_physical_table(result: std::result::Result<(), Arc>) { - let err = result.expect_err("the empty catalog should make the flush fail"); - assert!( - matches!( - err.as_ref(), - Error::Internal { err_msg } - if err_msg.contains("not found during pending flush") - ), - "unexpected flush error: {err}" - ); - } - #[tokio::test(start_paused = true)] async fn test_worker_preserves_first_deadline_and_inline_shutdown() { let flush_interval = Duration::from_secs(10); @@ -562,7 +503,8 @@ mod tests { datanodes: Arc::new(HashMap::new()), }); let catalog_manager = MemoryCatalogManager::with_default_setup(); - let (flow_notification_tx, _flow_notification_rx) = Notifier::try_new(1).unwrap(); + let (flow_notification_tx, _flow_notification_rx) = + FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); let (shutdown, _) = broadcast::channel(1); let flush_limiter = FlushLimiter::try_new(1).unwrap(); @@ -653,4 +595,64 @@ mod tests { "worker did not exit after shutdown" ); } + + const WORKER_TEST_TIMEOUT: Duration = Duration::from_secs(30); + + async fn submit_mock_worker_batch( + worker_tx: &mpsc::Sender, + total_rows: usize, + timestamp: i64, + ) -> oneshot::Receiver>> { + let (response_tx, response_rx) = oneshot::channel(); + let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); + worker_tx + .send(WorkerCommand::Submit { + table_batches: vec![( + "cpu".to_string(), + 42, + mock_aligned_tag_batch("tag1", "host-1", timestamp, 1.0), + )], + total_rows, + ctx: session::context::QueryContext::arc(), + response_tx, + _permit: Arc::new(permit), + }) + .await + .unwrap(); + + // The channel is FIFO, so the ack proves the worker has dequeued and + // processed the submission (anchoring the flush deadline) before the + // caller advances virtual time. + let (ack_tx, ack_rx) = oneshot::channel(); + worker_tx.send(WorkerCommand::Ack { ack_tx }).await.unwrap(); + ack_rx + .await + .expect("worker exited before acking the submitted batch"); + + response_rx + } + + async fn receive_mock_flush_result( + response_rx: oneshot::Receiver>>, + context: &str, + ) -> std::result::Result<(), Arc> { + // Under paused time the timeout auto-advances the clock and fires + // deterministically if the flush never completes. + tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx) + .await + .unwrap_or_else(|_| panic!("{context}")) + .expect("flush result channel closed without a result") + } + + fn assert_missing_physical_table(result: std::result::Result<(), Arc>) { + let err = result.expect_err("the empty catalog should make the flush fail"); + assert!( + matches!( + err.as_ref(), + Error::Internal { err_msg } + if err_msg.contains("not found during pending flush") + ), + "unexpected flush error: {err}" + ); + } } diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs index e3f817e70f2..71b558d1428 100644 --- a/src/servers/src/batcher/logical_table/region_write.rs +++ b/src/servers/src/batcher/logical_table/region_write.rs @@ -385,80 +385,6 @@ mod tests { }; use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager; - struct TwoRegionPartitionRule { - partition_columns: Vec, - } - - impl PartitionRule for TwoRegionPartitionRule { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn partition_columns(&self) -> &[String] { - &self.partition_columns - } - - fn find_region( - &self, - _values: &[datatypes::prelude::Value], - ) -> partition::error::Result { - unimplemented!() - } - - fn split_record_batch( - &self, - _record_batch: &RecordBatch, - ) -> partition::error::Result> - { - Ok(HashMap::from([ - (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), - (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), - ( - 3, - RegionMask::new(BooleanArray::from(vec![false, false]), 0), - ), - ])) - } - } - - #[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!() - } - } - #[tokio::test] async fn test_flush_region_writes_concurrently_dispatches_multiple_datanodes() { let inflight = Arc::new(AtomicUsize::new(0)); @@ -558,6 +484,80 @@ mod tests { check_encode_region_write_requests(true); } + struct TwoRegionPartitionRule { + partition_columns: Vec, + } + + impl PartitionRule for TwoRegionPartitionRule { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn partition_columns(&self) -> &[String] { + &self.partition_columns + } + + fn find_region( + &self, + _values: &[datatypes::prelude::Value], + ) -> partition::error::Result { + unimplemented!() + } + + fn split_record_batch( + &self, + _record_batch: &RecordBatch, + ) -> partition::error::Result> + { + Ok(HashMap::from([ + (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)), + (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)), + ( + 3, + RegionMask::new(BooleanArray::from(vec![false, false]), 0), + ), + ])) + } + } + + #[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), diff --git a/src/servers/src/batcher/logical_table/tables.rs b/src/servers/src/batcher/logical_table/tables.rs new file mode 100644 index 00000000000..4ed07f73a6e --- /dev/null +++ b/src/servers/src/batcher/logical_table/tables.rs @@ -0,0 +1,551 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use api::v1::{ColumnSchema, RowInsertRequests, Rows}; +use arrow::datatypes::Schema as ArrowSchema; +use async_trait::async_trait; +use session::context::QueryContextRef; +use snafu::OptionExt; + +use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; +use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx; +use crate::error; +use crate::error::Result; +use crate::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED; +use crate::prom_row_builder::{ + build_prom_create_table_schema_from_proto, identify_missing_columns_from_proto, + rows_to_aligned_record_batch, +}; + +#[async_trait] +pub trait PendingRowsSchemaAlterer: Send + Sync { + /// Batch-create multiple logical tables that are missing. + /// Each entry is `(table_name, request_schema)`. + async fn create_tables_if_missing_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[ColumnSchema])], + with_metric_engine: bool, + ctx: QueryContextRef, + ) -> Result<()>; + + /// Batch-alter multiple logical tables to add missing tag columns. + /// Each entry is `(table_name, missing_column_names)`. + async fn add_missing_prom_tag_columns_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[String])], + ctx: QueryContextRef, + ) -> Result<()>; +} + +pub type PendingRowsSchemaAltererRef = Arc; + +/// Intermediate planning state for resolving and preparing logical tables +/// before row-to-batch alignment. +pub(in crate::batcher::logical_table) struct TableResolutionPlan { + /// Resolved table schema and table id by logical table name. + pub(in crate::batcher::logical_table) region_schemas: HashMap, u32)>, + /// Missing tables that need to be created before alignment. + pub(in crate::batcher::logical_table) tables_to_create: Vec<(String, Vec)>, + /// Existing tables that need tag-column schema evolution. + pub(in crate::batcher::logical_table) tables_to_alter: Vec<(String, Vec)>, +} + +impl LogicalTablePendingRowsBatcher { + /// Converts proto `RowInsertRequests` directly into aligned `RecordBatch`es + /// in a single pass, handling table creation, schema alteration, column + /// renaming, reordering, and null-filling without building intermediate + /// RecordBatches. + pub(in crate::batcher::logical_table) async fn build_and_align_table_batches( + &self, + requests: RowInsertRequests, + ctx: &QueryContextRef, + ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { + let catalog = ctx.current_catalog().to_string(); + let schema = ctx.current_schema(); + + let (table_rows, total_rows) = Self::collect_non_empty_table_rows(requests); + if total_rows == 0 { + return Ok((Vec::new(), 0)); + } + + let unique_tables = Self::collect_unique_table_schemas(&table_rows)?; + let mut plan = self + .plan_table_resolution(&catalog, &schema, ctx, &unique_tables) + .await?; + + self.create_missing_tables_and_refresh_schemas( + &catalog, + &schema, + ctx, + &table_rows, + &mut plan, + ) + .await?; + + self.alter_tables_and_refresh_schemas(&catalog, &schema, ctx, &mut plan) + .await?; + + let aligned_batches = Self::build_aligned_batches(&table_rows, &plan.region_schemas)?; + + Ok((aligned_batches, total_rows)) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Extracts non-empty `(table_name, rows)` pairs and computes total row + /// count across the retained entries. + pub(in crate::batcher::logical_table) fn collect_non_empty_table_rows( + requests: RowInsertRequests, + ) -> (Vec<(String, Rows)>, usize) { + let mut table_rows: Vec<(String, Rows)> = Vec::with_capacity(requests.inserts.len()); + let mut total_rows = 0; + + for request in requests.inserts { + let Some(rows) = request.rows else { + continue; + }; + if rows.rows.is_empty() { + continue; + } + + total_rows += rows.rows.len(); + table_rows.push((request.table_name, rows)); + } + + (table_rows, total_rows) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Returns unique `(table_name, proto_schema)` pairs while keeping the + /// first-seen schema for duplicate table names. + pub(in crate::batcher::logical_table) fn collect_unique_table_schemas( + table_rows: &[(String, Rows)], + ) -> Result> { + let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); + let mut seen = HashSet::new(); + + for (table_name, rows) in table_rows { + if seen.insert(table_name.as_str()) { + unique_tables.push((table_name.as_str(), &rows.schema)); + } else { + // table_rows should group rows by table name. + return error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Found duplicated table name in RowInsertRequest: {}", + table_name + ), + } + .fail(); + } + } + + Ok(unique_tables) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Resolves table metadata and classifies each table into existing, + /// to-create, and to-alter groups used by subsequent DDL steps. + pub(in crate::batcher::logical_table) async fn plan_table_resolution( + &self, + catalog: &str, + schema: &str, + ctx: &QueryContextRef, + unique_tables: &[(&str, &[ColumnSchema])], + ) -> Result { + let mut plan = TableResolutionPlan { + region_schemas: HashMap::with_capacity(unique_tables.len()), + tables_to_create: Vec::new(), + tables_to_alter: Vec::new(), + }; + + let resolved_tables = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_resolve_table"]) + .start_timer(); + futures::future::join_all(unique_tables.iter().map(|(table_name, _)| { + self.catalog_manager + .table(catalog, schema, table_name, Some(ctx.as_ref())) + })) + .await + }; + + for ((table_name, rows_schema), table_result) in unique_tables.iter().zip(resolved_tables) { + let table = table_result?; + + if let Some(table) = table { + let table_info = table.table_info(); + let table_id = table_info.ident.table_id; + let region_schema = table_info.meta.schema.arrow_schema().clone(); + + let missing_columns = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_identify_missing_columns"]) + .start_timer(); + identify_missing_columns_from_proto(rows_schema, region_schema.as_ref())? + }; + if !missing_columns.is_empty() { + plan.tables_to_alter + .push(((*table_name).to_string(), missing_columns)); + } + plan.region_schemas + .insert((*table_name).to_string(), (region_schema, table_id)); + } else { + let request_schema = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_build_create_table_schema"]) + .start_timer(); + build_prom_create_table_schema_from_proto(rows_schema)? + }; + plan.tables_to_create + .push(((*table_name).to_string(), request_schema)); + } + } + + Ok(plan) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Batch-creates missing tables, refreshes their schema metadata, and + /// enqueues follow-up alters for extra tag columns discovered in later rows. + pub(in crate::batcher::logical_table) async fn create_missing_tables_and_refresh_schemas( + &self, + catalog: &str, + schema: &str, + ctx: &QueryContextRef, + table_rows: &[(String, Rows)], + plan: &mut TableResolutionPlan, + ) -> Result<()> { + if plan.tables_to_create.is_empty() { + return Ok(()); + } + + let create_refs: Vec<(&str, &[ColumnSchema])> = plan + .tables_to_create + .iter() + .map(|(name, schema)| (name.as_str(), schema.as_slice())) + .collect(); + + { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_batch_create_tables"]) + .start_timer(); + self.schema_alterer + .create_tables_if_missing_batch( + catalog, + schema, + &create_refs, + self.prom_store_with_metric_engine, + ctx.clone(), + ) + .await?; + } + + let created_table_results = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_resolve_table_after_create"]) + .start_timer(); + futures::future::join_all(plan.tables_to_create.iter().map(|(table_name, _)| { + self.catalog_manager + .table(catalog, schema, table_name, Some(ctx.as_ref())) + })) + .await + }; + + for ((table_name, _), table_result) in + plan.tables_to_create.iter().zip(created_table_results) + { + let table = table_result?.with_context(|| error::UnexpectedResultSnafu { + reason: format!( + "Table not found after pending batch create attempt: {}", + table_name + ), + })?; + let table_info = table.table_info(); + let table_id = table_info.ident.table_id; + let region_schema = table_info.meta.schema.arrow_schema().clone(); + plan.region_schemas + .insert(table_name.clone(), (region_schema, table_id)); + } + + Self::enqueue_alter_for_new_tables(table_rows, plan)?; + + Ok(()) + } +} + +impl LogicalTablePendingRowsBatcher { + /// For newly created tables, re-checks all row schemas and appends alter + /// operations when additional tag columns are still missing. + pub(in crate::batcher::logical_table) fn enqueue_alter_for_new_tables( + table_rows: &[(String, Rows)], + plan: &mut TableResolutionPlan, + ) -> Result<()> { + let created_tables: HashSet<&str> = plan + .tables_to_create + .iter() + .map(|(table_name, _)| table_name.as_str()) + .collect(); + + for (table_name, rows) in table_rows { + if !created_tables.contains(table_name.as_str()) { + continue; + } + + let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { + continue; + }; + + let missing_columns = identify_missing_columns_from_proto(&rows.schema, region_schema)?; + if missing_columns.is_empty() + || plan + .tables_to_alter + .iter() + .any(|(existing_name, _)| existing_name == table_name) + { + continue; + } + + plan.tables_to_alter + .push((table_name.clone(), missing_columns)); + } + + Ok(()) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Batch-alters tables that have missing tag columns and refreshes the + /// in-memory schema map used for row alignment. + pub(in crate::batcher::logical_table) async fn alter_tables_and_refresh_schemas( + &self, + catalog: &str, + schema: &str, + ctx: &QueryContextRef, + plan: &mut TableResolutionPlan, + ) -> Result<()> { + if plan.tables_to_alter.is_empty() { + return Ok(()); + } + + let alter_refs: Vec<(&str, &[String])> = plan + .tables_to_alter + .iter() + .map(|(name, cols)| (name.as_str(), cols.as_slice())) + .collect(); + { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_batch_add_missing_columns"]) + .start_timer(); + self.schema_alterer + .add_missing_prom_tag_columns_batch(catalog, schema, &alter_refs, ctx.clone()) + .await?; + } + + let altered_table_results = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_resolve_table_after_schema_alter"]) + .start_timer(); + futures::future::join_all(plan.tables_to_alter.iter().map(|(table_name, _)| { + self.catalog_manager + .table(catalog, schema, table_name, Some(ctx.as_ref())) + })) + .await + }; + + for ((table_name, _), table_result) in + plan.tables_to_alter.iter().zip(altered_table_results) + { + let table = table_result?.with_context(|| error::UnexpectedResultSnafu { + reason: format!( + "Table not found after pending batch schema alter: {}", + table_name + ), + })?; + let table_info = table.table_info(); + let table_id = table_info.ident.table_id; + let refreshed_region_schema = table_info.meta.schema.arrow_schema().clone(); + plan.region_schemas + .insert(table_name.clone(), (refreshed_region_schema, table_id)); + } + + Ok(()) + } +} + +impl LogicalTablePendingRowsBatcher { + /// Converts proto rows to `RecordBatch` values aligned to resolved region + /// schemas and returns `(table_name, table_id, batch)` tuples. + pub(in crate::batcher::logical_table) fn build_aligned_batches( + table_rows: &[(String, Rows)], + region_schemas: &HashMap, u32)>, + ) -> Result> { + let mut aligned_batches = Vec::with_capacity(table_rows.len()); + for (table_name, rows) in table_rows { + let (region_schema, table_id) = + region_schemas.get(table_name).cloned().with_context(|| { + error::UnexpectedResultSnafu { + reason: format!("Region schema not resolved for table: {}", table_name), + } + })?; + + let record_batch = { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["align_rows_to_record_batch"]) + .start_timer(); + rows_to_aligned_record_batch(rows, region_schema.as_ref())? + }; + aligned_batches.push((table_name.clone(), table_id, record_batch)); + } + + Ok(aligned_batches) + } +} + +#[cfg(test)] +mod tests { + + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, + Value, + }; + use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; + use common_query::prelude::greptime_timestamp; + + use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; + use crate::batcher::logical_table::batch_convert::TableBatch; + use crate::batcher::logical_table::flow_notifier::extract_timestamps; + use crate::prom_row_builder::rows_to_aligned_record_batch; + + #[test] + fn test_extract_timestamps_uses_aligned_custom_timestamp_index() { + let rows = Rows { + schema: vec![ + ColumnSchema { + column_name: greptime_timestamp().to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "host".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ], + rows: vec![ + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }, + Value { + value_data: Some(ValueData::StringValue("host-1".to_string())), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }, + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(2000)), + }, + Value { + value_data: Some(ValueData::StringValue("host-2".to_string())), + }, + Value { + value_data: Some(ValueData::F64Value(2.0)), + }, + ], + }, + ], + }; + let target_schema = ArrowSchema::new(vec![ + Field::new("host", ArrowDataType::Utf8, true), + Field::new( + "timestamp", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", ArrowDataType::Float64, true), + ]); + let batch = rows_to_aligned_record_batch(&rows, &target_schema).unwrap(); + assert_eq!(1, batch.timestamp_index); + let table_batch = TableBatch { + table_name: "cpu".to_string(), + table_id: 42, + row_count: batch.batch.num_rows(), + batches: vec![batch], + }; + + assert_eq!(vec![1000, 2000], extract_timestamps(&table_batch)); + } + + #[test] + fn test_collect_non_empty_table_rows_filters_empty_payloads() { + let requests = RowInsertRequests { + inserts: vec![ + RowInsertRequest { + table_name: "cpu".to_string(), + rows: Some(mock_rows(2, "host")), + }, + RowInsertRequest { + table_name: "mem".to_string(), + rows: Some(mock_rows(0, "host")), + }, + RowInsertRequest { + table_name: "disk".to_string(), + rows: None, + }, + ], + }; + + let (table_rows, total_rows) = + LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); + + assert_eq!(2, total_rows); + assert_eq!(1, table_rows.len()); + assert_eq!("cpu", table_rows[0].0); + assert_eq!(2, table_rows[0].1.rows.len()); + } + + fn mock_rows(row_count: usize, schema_name: &str) -> Rows { + Rows { + schema: vec![ColumnSchema { + column_name: schema_name.to_string(), + ..Default::default() + }], + rows: (0..row_count).map(|_| Row { values: vec![] }).collect(), + } + } +} diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs index b04883b27e0..1cc81f9b615 100644 --- a/src/servers/src/batcher/logical_table/test_util.rs +++ b/src/servers/src/batcher/logical_table/test_util.rs @@ -23,16 +23,11 @@ use arrow::array::{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::cache::{TableFlownodeSetCacheRef, new_table_flownode_set_cache}; use common_meta::error::Result as MetaResult; -use common_meta::instruction::{CacheIdent, CreateFlow}; -use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::{ DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, }; -use moka::future::CacheBuilder; use snafu::ResultExt; -use table::metadata::TableId; use tokio::sync::mpsc; use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx; @@ -155,26 +150,6 @@ impl FlownodeManager for FlowNotificationMockNodeManager { } } -pub(in crate::batcher::logical_table) async fn mock_table_flownode_cache( - table_id: TableId, - peer: Peer, -) -> TableFlownodeSetCacheRef { - let cache = Arc::new(new_table_flownode_set_cache( - "test".to_string(), - CacheBuilder::new(1).build(), - Arc::new(MemoryKvBackend::default()), - )); - cache - .invalidate(&[CacheIdent::CreateFlow(CreateFlow { - flow_id: 1, - source_table_ids: vec![table_id], - partition_to_peer_mapping: vec![(0, peer.clone()), (1, peer)], - })]) - .await - .unwrap(); - cache -} - #[async_trait] impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager { async fn handle(&self, peer: &Peer, request: RegionRequest) -> error::Result { From 9dfe057199b75a9f3dfb9d16b1188db8b41d77e4 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 04:58:58 +0000 Subject: [PATCH 50/63] test: make export chunk deletion failure deterministic (#9291) * test: make export chunk deletion failure deterministic Signed-off-by: jeremyhi * test: simplify export chunk deletion failure fixture Signed-off-by: jeremyhi * docs: guide deterministic storage failure tests Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/cli/Cargo.toml | 1 + src/cli/src/data/snapshot_storage.rs | 37 ++++++++++++++++++---------- tests-integration/AGENTS.md | 4 +++ 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 3a7dc1d68e4..587e5b8cffa 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -76,5 +76,6 @@ uuid.workspace = true common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-version.workspace = true +object-store = { workspace = true, features = ["testing"] } serde.workspace = true tempfile.workspace = true diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index b3a5fca84bc..c5e9561d285 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -796,13 +796,16 @@ fn valid_chunk_filename(name: &str) -> bool { mod tests { use std::collections::HashMap; use std::path::Path; + use std::sync::Arc; use object_store::ObjectStore; + use object_store::layers::mock::MockLayerBuilder; use object_store::services::Fs; use tempfile::tempdir; use url::Url; use super::*; + use crate::data::export_v2::error::Error; use crate::data::export_v2::manifest::{DataFormat, TimeRange}; use crate::data::export_v2::schema::SchemaDefinition; @@ -1140,28 +1143,36 @@ mod tests { } } - #[cfg(unix)] #[tokio::test] async fn test_prepare_export_chunk_reports_delete_failure() { - use std::os::unix::fs::PermissionsExt; let dir = tempdir().unwrap(); - let storage = make_storage_with_rooted_fs(dir.path()); + let mut storage = make_storage_with_rooted_fs(dir.path()); + storage.object_store = storage.object_store.layer( + MockLayerBuilder::default() + // OpenDAL's unit deleter returns Unsupported for every deletion. + .deleter_factory(Arc::new(|_| Box::new(()))) + .build() + .unwrap(), + ); storage .write_text("data/public/2/a.parquet", "keep") .await .unwrap(); - let chunk_dir = dir.path().join("data/public/2"); - std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o555)).unwrap(); - let result = storage + let error = storage .prepare_export_chunk(&["public".into()], 2, true) - .await; - std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o755)).unwrap(); - assert!( - result - .unwrap_err() - .to_string() - .contains("delete unfinished chunk file") + .await + .unwrap_err(); + let Error::StorageOperation { + operation, error, .. + } = error + else { + panic!("expected storage operation error, got {error:?}"); + }; + assert_eq!( + operation, + "delete unfinished chunk file data/public/2/a.parquet" ); + assert_eq!(error.kind(), ErrorKind::Unsupported); assert_eq!( storage.read_text("data/public/2/a.parquet").await.unwrap(), "keep" diff --git a/tests-integration/AGENTS.md b/tests-integration/AGENTS.md index 9b83a05497c..25531aeb1b9 100644 --- a/tests-integration/AGENTS.md +++ b/tests-integration/AGENTS.md @@ -20,6 +20,10 @@ COPY locations. paths to file URIs with `Url::from_file_path`; do not hard-code `/tmp` or concatenate `file://` with native paths. Close file handles before rename or cleanup, including import-state writes. +- Linux [Nightly Build](../.github/workflows/nightly-build.yml) runs tests as root in + `dev-builder`; `chmod`/read-only permissions cannot reliably force I/O failures. + For storage error propagation tests, reuse `object_store::layers::mock` to inject + deterministic failures rather than relying on execution-user permissions. ## Validation From f1e9a74f003c94b7e7fdcfd2576da0e390891224 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Tue, 22 Sep 2026 07:35:07 +0000 Subject: [PATCH 51/63] fix(meta): populate physical metric table column ids (#9286) * fix(meta): populate physical metric table column ids Signed-off-by: dhruvxvaishnav * test(meta): verify physical metric column ids Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav --- .../meta/src/ddl/create_table/executor.rs | 14 ++- tests-integration/tests/repartition.rs | 95 +++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) 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/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index c243113d928..5fe84b19f24 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() { From 614b46e5f52586ad54f617822ff3f026491ee88e Mon Sep 17 00:00:00 2001 From: houyuwushang <180804215+houyuwushang@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:03:40 +0000 Subject: [PATCH 52/63] fix(postgres): return empty responses for comment-only SQL (#9295) fix(postgres): handle parsed empty queries in both protocols Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com> --- src/servers/src/postgres/handler.rs | 7 +++ src/servers/tests/postgres/mod.rs | 86 +++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 204fa7745a5..e618c55614b 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -82,6 +82,10 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { let parsed_query = self.query_parser.compatibility_parser.parse(query); let query = if let Ok(statements) = &parsed_query { + // Comments, whitespace and empty statements also require EmptyQueryResponse. + if statements.is_empty() { + return Ok(vec![Response::EmptyQuery]); + } statements .iter() .map(|s| s.to_string()) @@ -335,6 +339,9 @@ impl QueryParser for DefaultQueryParser { let parsed_statements = self.compatibility_parser.parse(sql); let (sql, copy_to_stdout_format) = if let Ok(mut statements) = parsed_statements { + if statements.is_empty() { + return Ok(None); + } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); (first_stmt.to_string(), format) diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 2e031c0d0db..6b2995283f8 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -587,6 +587,92 @@ async fn start_test_server(server_tls: TlsOption) -> Result { Ok(server_addr.port()) } +const EMPTY_QUERIES: &[&str] = &[ + "", + " \t\r\n", + ";", + ";;;", + "-- ping", + "-- ping\n", + "/* ping */", + "/* outer /* inner */ comment */", + "; -- ping\r\n /* comment */ ;", +]; + +#[tokio::test] +async fn test_simple_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let messages = client + .simple_query(query) + .await + .unwrap_or_else(|err| panic!("query {query:?} failed: {err}")); + // tokio-postgres exposes EmptyQueryResponse as CommandComplete(0). + assert!( + matches!( + messages.as_slice(), + [SimpleQueryMessage::CommandComplete(0)] + ), + "expected an empty query response for {query:?}, got {messages:?}" + ); + } + + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + +#[tokio::test] +async fn test_extended_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let statement = client + .prepare(query) + .await + .unwrap_or_else(|err| panic!("prepare {query:?} failed: {err}")); + assert_eq!(client.execute(&statement, &[]).await.unwrap(), 0); + assert!(client.query(&statement, &[]).await.unwrap().is_empty()); + } + + let row = client + .query_one("SELECT $1::BIGINT", &[&42i64]) + .await + .unwrap(); + assert_eq!(row.get::<_, i64>(0), 42); + Ok(()) +} + +#[tokio::test] +async fn test_simple_query_with_comments() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for (query, expected) in [ + ("-- ping\nSELECT 1", vec!["1"]), + ("/* comment */ SELECT 1; -- trailing", vec!["1"]), + ("SELECT '-- ping'", vec!["-- ping"]), + ("SELECT '/* comment */'", vec!["/* comment */"]), + ("SELECT 1; /* between */ SELECT 2;", vec!["1", "2"]), + ] { + let messages = client.simple_query(query).await.unwrap(); + assert_eq!(unwrap_results(&messages), expected, "query: {query}"); + } + + assert!( + client + .simple_query("/* comment */ SELECT missing_column FROM numbers") + .await + .is_err() + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + async fn do_simple_query(server_tls: TlsOption, client_tls: bool) -> Result<()> { let server_port = start_test_server(server_tls).await?; From 2388257c35beb16e50132d35b077ae26a4183607 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 22 Sep 2026 08:37:21 +0000 Subject: [PATCH 53/63] ci: create docs follow-up issue on PR merge instead of on label (#9237) * ci: create docs follow-up issue on PR merge instead of on label The docbot workflow previously created a docs-repo issue as soon as the 'docs-required' condition was detected (PR opened/edited with the docs checkbox ticked), even if the PR was never merged. Now the workflow also triggers on PR 'closed': - opened/edited: only manage the docs-required/docs-not-required labels - closed: create the docs issue only when the PR was actually merged and carries the docs-required label This also lets maintainers control issue creation by manually adding or removing the docs-required label before merging. Signed-off-by: Ning Sun * fix: address review comments on docs issue creation timing - Only touch docs labels when the docs checkbox state actually changed in an edit. Previously, editing any other part of the PR body while the checkbox stayed checked removed the docs-required label, silently dropping the docs follow-up now that issue creation happens at merge. Unchanged checkbox now leaves labels untouched, which also preserves manual label overrides. - Do not trust the closed event's stale label snapshot at merge time: re-read the live PR via the API and create the docs issue if the docs-required label is present OR the checkbox is ticked in the current body. - Make the workflow concurrency group action-aware so a merge run does not cancel an in-flight label update from an edit run. Signed-off-by: Ning Sun * fix: make docs-required label the single source of truth at merge The label-OR-checkbox merge condition could not distinguish an intentional opt-out from an unfinished label update: removing docs-required while the checkbox stayed checked still produced an issue, and unchecking the box could still produce one if the merge read the stale label before the edit run removed it. At merge time, wait for any pending docbot runs on the PR head SHA to finish their label updates (bounded to 5 minutes), then decide solely by the live docs-required label. Adds actions: read permission for listing workflow runs. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .github/workflows/docbot.yml | 5 +- cyborg/bin/follow-up-docs-issue.ts | 153 +++++++++++++++++++++-------- 2 files changed, 114 insertions(+), 44 deletions(-) diff --git a/.github/workflows/docbot.yml b/.github/workflows/docbot.yml index 250a150c026..53033c6eb5b 100644 --- a/.github/workflows/docbot.yml +++ b/.github/workflows/docbot.yml @@ -1,10 +1,10 @@ name: Follow Up Docs on: pull_request_target: - types: [opened, edited] + types: [opened, edited, closed] concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + group: ${{ github.workflow }}-${{ github.event.action }}-${{ github.head_ref || github.run_id }} cancel-in-progress: true jobs: @@ -13,6 +13,7 @@ jobs: permissions: pull-requests: write contents: read + actions: read timeout-minutes: 10 steps: - uses: actions/checkout@v4 diff --git a/cyborg/bin/follow-up-docs-issue.ts b/cyborg/bin/follow-up-docs-issue.ts index acca21fe935..09ef5db2445 100644 --- a/cyborg/bin/follow-up-docs-issue.ts +++ b/cyborg/bin/follow-up-docs-issue.ts @@ -17,7 +17,12 @@ import * as core from '@actions/core' import {handleError, obtainClient} from "@/common"; import {context} from "@actions/github"; -import {PullRequestEditedEvent, PullRequestEvent, PullRequestOpenedEvent} from "@octokit/webhooks-types"; +import { + PullRequestClosedEvent, + PullRequestEditedEvent, + PullRequestEvent, + PullRequestOpenedEvent +} from "@octokit/webhooks-types"; // @ts-expect-error moduleResolution:nodenext issue 54523 import {RequestError} from "@octokit/request-error"; @@ -30,18 +35,59 @@ async function main() { throw new Error(`Only pull request event supported. ${context.eventName} is unsupported.`) } - const client = obtainClient("GITHUB_TOKEN") - const docsClient = obtainClient("DOCS_REPO_TOKEN") const payload = context.payload as PullRequestEvent - const { owner, repo, number, actor, title, html_url } = { + const { owner, repo, number, actor } = { owner: payload.pull_request.base.user.login, repo: payload.pull_request.base.repo.name, number: payload.pull_request.number, - title: payload.pull_request.title, - html_url: payload.pull_request.html_url, actor: payload.pull_request.user.login, } - const followUpDocs = checkPullRequestEvent(payload) + + switch (payload.action) { + case "opened": { + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent)) + break + } + case "edited": { + const followUpDocs = checkPullRequestEditedEvent(payload as PullRequestEditedEvent) + if (followUpDocs === undefined) { + core.info("Docs checkbox state unchanged; leaving labels as-is.") + return + } + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, followUpDocs) + break + } + case "closed": { + const event = payload as PullRequestClosedEvent + if (!event.pull_request.merged) { + core.info("PR closed without merging; no docs issue needed.") + return + } + // The docs-required label is the single source of truth. Coordinate + // with pending checkbox-driven label updates before deciding: an edit + // run may still be applying or removing the label, and the event + // payload's label snapshot may be stale. + const client = obtainClient("GITHUB_TOKEN") + const { data: pr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + await waitForPendingLabelUpdates(client, owner, repo, pr.head.sha) + const { data: freshPr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + const hasDocsLabel = freshPr.labels.some((label) => label.name === labelDocsRequired) + if (!hasDocsLabel) { + core.info(`Label ${labelDocsRequired} not present; no docs issue needed.`) + return + } + const docsClient = obtainClient("DOCS_REPO_TOKEN") + await createDocsIssue(docsClient, freshPr.title, freshPr.html_url, actor) + break + } + default: + throw new Error(`${payload.action} is unsupported.`) + } +} + +async function updateDocsLabels(client: ReturnType, owner: string, repo: string, number: number, followUpDocs: boolean) { if (followUpDocs) { core.info("Follow up docs.") await client.rest.issues.removeLabel({ @@ -55,28 +101,6 @@ async function main() { await client.rest.issues.addLabels({ owner, repo, issue_number: number, labels: [labelDocsRequired], }) - - // Get available assignees for the docs repo - const assigneesResponse = await docsClient.rest.issues.listAssignees({ - owner: 'GreptimeTeam', - repo: 'docs', - }) - const validAssignees = assigneesResponse.data.map(assignee => assignee.login) - core.info(`Available assignees: ${validAssignees.join(', ')}`) - - // Check if the actor is a valid assignee, otherwise fallback to fengjiachun - const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' - core.info(`Assigning issue to: ${assignee}`) - - await docsClient.rest.issues.create({ - owner: 'GreptimeTeam', - repo: 'docs', - title: `Update docs for ${title}`, - body: `A document change request is generated from ${html_url}`, - assignee: assignee, - }).then((res) => { - core.info(`Created issue ${res.data}`) - }) } else { core.info("No need to follow up docs.") await client.rest.issues.removeLabel({ @@ -93,15 +117,52 @@ async function main() { } } -function checkPullRequestEvent(payload: PullRequestEvent) { - switch (payload.action) { - case "opened": - return checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent) - case "edited": - return checkPullRequestEditedEvent(payload as PullRequestEditedEvent) - default: - throw new Error(`${payload.action} is unsupported.`) +async function createDocsIssue(docsClient: ReturnType, title: string, html_url: string, actor: string) { + core.info("Creating follow-up docs issue for merged PR.") + + // Get available assignees for the docs repo + const assigneesResponse = await docsClient.rest.issues.listAssignees({ + owner: 'GreptimeTeam', + repo: 'docs', + }) + const validAssignees = assigneesResponse.data.map(assignee => assignee.login) + core.info(`Available assignees: ${validAssignees.join(', ')}`) + + // Check if the actor is a valid assignee, otherwise fallback to fengjiachun + const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' + core.info(`Assigning issue to: ${assignee}`) + + await docsClient.rest.issues.create({ + owner: 'GreptimeTeam', + repo: 'docs', + title: `Update docs for ${title}`, + body: `A document change request is generated from ${html_url}`, + assignee: assignee, + }).then((res) => { + core.info(`Created issue ${res.data}`) + }) +} + +// Waits until no other runs of this workflow are pending for the PR head SHA, +// so that any checkbox-driven label update has been applied before the caller +// reads the labels. Bounded by a deadline; on timeout it proceeds with whatever +// the current label state is. +async function waitForPendingLabelUpdates(client: ReturnType, owner: string, repo: string, headSha: string) { + const deadlineMs = 5 * 60 * 1000 + const pollIntervalMs = 5000 + const deadline = Date.now() + deadlineMs + while (Date.now() < deadline) { + const { data } = await client.rest.actions.listWorkflowRuns({ + owner, repo, workflow_id: "docbot.yml", head_sha: headSha, per_page: 20, + }) + const pending = data.workflow_runs.filter((run) => run.id !== context.runId && run.status !== "completed") + if (pending.length === 0) { + return + } + core.info(`Waiting for ${pending.length} pending docbot run(s) to finish label updates...`) + await new Promise((resolve) => setTimeout(resolve, pollIntervalMs)) } + core.warning("Timed out waiting for pending docbot runs; proceeding with the current label state.") } function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { @@ -109,11 +170,19 @@ function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { return event.pull_request.body?.includes(needFollowUpDocs) } -function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean { - const previous = event.changes.body?.from.includes(needFollowUpDocs) - const current = event.pull_request.body?.includes(needFollowUpDocs) - // from docs-not-need to docs-required - return (!previous) && current +// Returns undefined when the checkbox state did not change in this edit, so the +// caller leaves the labels untouched (preserving manual label overrides). +function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean | undefined { + if (!event.changes.body) { + // The body was not part of this edit (e.g. title-only edit). + return undefined + } + const previous = event.changes.body.from.includes(needFollowUpDocs) + const current = event.pull_request.body?.includes(needFollowUpDocs) ?? false + if (previous === current) { + return undefined + } + return current } main().catch(handleError) From 25bca4609f99fdd46078c528e4ad84aea7f7281f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 11:25:50 +0000 Subject: [PATCH 54/63] perf(promql): push label filters into grouped join inputs (#9280) * perf(promql): propagate matching filters through grouped joins Signed-off-by: Dennis Zhuang * perf(promql): check matcher safety on the receiving operand Signed-off-by: Dennis Zhuang * refactor(promql): spell out the shapes a filter may cross `preserves_filter` ended in `_ => true`, which was only sound because `selector_matchers` independently rejects label rewriting, `count_values`, subqueries and non-rollup calls on the same operand. Loosening the latter alone would have silently pushed a matcher below a label rewrite. List the shapes that carry a scan filter instead and default to `false`. Cite #9207 for the result labels the grouped cases record: the join projects the right operand's tag set, so `zone` is missing wherever the right side aggregates it away. No behavior change. Signed-off-by: Dennis Zhuang * test(promql): assert the new pushdowns reach the scan The grouped-join unit tests feed tag columns by hand and the SQLness case only checks results, which are identical whether or not the rewrite fires. Nothing would have failed if scalar arithmetic, ranking or grouped matching stopped propagating. Assert through the planner that the matcher reaches both scans, with a global topk one-side as the counter-example. Also state that the duplicate-one-side cases record a cross product Prometheus rejects (#9209), so the baseline is not read as intended semantics. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/planner.rs | 43 ++- .../src/promql/planner/matching_filters.rs | 300 ++++++++++++------ .../common/promql/matching_groups.result | 181 +++++++++++ .../common/promql/matching_groups.sql | 65 ++++ 4 files changed, 484 insertions(+), 105 deletions(-) create mode 100644 tests/cases/standalone/common/promql/matching_groups.result create mode 100644 tests/cases/standalone/common/promql/matching_groups.sql diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 94081c2d284..44211b8d2b7 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -13515,6 +13515,30 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { + for query in [ + r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, + r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, + r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } + // A global ranking one-side must see every host, so the matcher stays put. + let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 1, + "{query}\n{plan}" + ); + } + #[tokio::test] async fn binary_matching_label_filter_skips_selecting_aggregations() { // `topk` ranks its input, so filtering before it changes the candidate set. @@ -13539,13 +13563,18 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom async fn binary_value_field_matcher_is_not_copied_across_aggregations() { // `status` varies between the samples of one series, so filtering the other operand by it // would drop the newest sample before sample selection (#9242). - let query = r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#; - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!( - plan.matches(r#"Utf8("ready")"#).count(), - 1, - "{query}\n{plan}" - ); + for query in [ + r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, + r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + ] { + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!( + plan.matches(r#"Utf8("ready")"#).count(), + 1, + "{query}\n{plan}" + ); + } } #[tokio::test] diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs index 386d95e9228..e38d48044da 100644 --- a/src/query/src/promql/planner/matching_filters.rs +++ b/src/query/src/promql/planner/matching_filters.rs @@ -52,26 +52,23 @@ const LABEL_PRESERVING_RANGE_FUNCTIONS: [&str; 20] = [ /// Copies label matchers from one operand of `binary` to the other, so both scans discard /// non-joining series instead of the join. /// -/// One-to-one arithmetic operands are inner-joined on their matching labels with plain column +/// Arithmetic operands are inner-joined on their matching labels with plain column /// equality, so a matcher on a matching label already holds for every surviving pair. Copying it /// to the other operand can only drop rows that had no partner, whatever the matcher kind, and /// `NULL` pairs only with `NULL`. Nor can it split a match group, whose series agree on every -/// matching label, so it does not affect a cardinality check such as #9209. +/// matching label, so it does not affect a one-to-one cardinality check such as #9209. +/// Grouped matching is conservatively restricted to a provably unique one-side operand until +/// its cardinality checks are implemented (#9208). /// /// `left_tags` and `right_tags` are the tag columns of the planned operands: a matcher may just /// as well constrain a value field, which the parsed expression does not distinguish from a -/// label. For an aggregated operand they are its grouping labels, which is what makes filtering -/// its input equivalent to filtering its output. +/// label. Matchers can cross an aggregate only on labels that partition its input. +/// `left_field_labels` and `right_field_labels` exclude value fields promoted to grouping +/// labels: filtering them before sample selection could replace the latest sample with an +/// older one (#9242). /// -/// A value field is not a series identity though: its value may vary between the samples of one -/// series, so a matcher on it must not be lowered below sample selection (`PromInstantManipulate`, -/// `PromSeriesDivide`), where discarding a sample promotes an older one out of the lookback window -/// and fabricates a match. An aggregated operand's tag columns are its grouping labels rather than -/// the tags of its input, so those labels may name value fields; a matcher therefore only crosses -/// an aggregation when it names one of the grouping labels that aggregation partitions by, and -/// `left_field_labels`/`right_field_labels` rule out the grouping labels the planner found to be -/// value fields of the aggregated operand's input -/// (`PromPlannerContext::aggregation_field_labels`). +/// Whether a matcher reaches a scan is decided per receiving operand ([`preserves_filter`]), so +/// propagation is one-way when only one of the two can take the filter. pub(super) fn propagate( binary: &BinaryExpr, left_tags: &[String], @@ -108,35 +105,16 @@ fn try_propagate( ) { return Err("operator is not arithmetic"); } - if let Some(modifier) = &binary.modifier { - if !matches!(modifier.card, VectorMatchCardinality::OneToOne) { - return Err("matching is not one-to-one"); - } - if modifier.fill_values.lhs.is_some() || modifier.fill_values.rhs.is_some() { - return Err("operand carries a fill modifier"); - } + if let Some(modifier) = &binary.modifier + && (modifier.fill_values.lhs.is_some() || modifier.fill_values.rhs.is_some()) + { + return Err("operand carries a fill modifier"); } let matching = binary .modifier .as_ref() .and_then(|modifier| modifier.matching.as_ref()); - // The grouping labels of every aggregation the operands cross, collected before `rewritten` - // takes the mutable borrows that rule out borrowing `binary` again. - let left_grouped = modifier_label_names(&binary.lhs); - let left_ignored = modifier_excluded_label_names(&binary.lhs); - let right_grouped = modifier_label_names(&binary.rhs); - let right_ignored = modifier_excluded_label_names(&binary.rhs); - - let mut rewritten = binary.clone(); - let (left, left_restricted) = - targeted_selector(&mut rewritten.lhs).ok_or("left operand is not a selector")?; - let (right, right_restricted) = - targeted_selector(&mut rewritten.rhs).ok_or("right operand is not a selector")?; - if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { - return Err("selector has an or matcher group"); - } - let is_matching_label = |name: &String| { // PromQL reserves the `__` prefix; none of those names is a join key. !name.starts_with("__") @@ -150,9 +128,32 @@ fn try_propagate( Some(LabelModifier::Include(on)) => on.labels.contains(name), Some(LabelModifier::Exclude(ignoring)) => !ignoring.labels.contains(name), } - && (!left_restricted || crosses_aggregation(name, &left_grouped, &left_ignored)) - && (!right_restricted || crosses_aggregation(name, &right_grouped, &right_ignored)) }; + if let Some(modifier) = &binary.modifier { + let unique_on_matching_labels = |expr: &Expr, tags: &[String]| { + has_unique_aggregate_output(expr) && tags.iter().all(&is_matching_label) + }; + let supported_cardinality = match &modifier.card { + VectorMatchCardinality::OneToOne => true, + VectorMatchCardinality::ManyToOne(_) => { + unique_on_matching_labels(&binary.rhs, right_tags) + } + VectorMatchCardinality::OneToMany(_) => { + unique_on_matching_labels(&binary.lhs, left_tags) + } + VectorMatchCardinality::ManyToMany => false, + }; + if !supported_cardinality { + return Err("one-side uniqueness is not proven"); + } + } + let mut rewritten = binary.clone(); + let left = selector_matchers(&mut rewritten.lhs).ok_or("left operand is not a selector")?; + let right = selector_matchers(&mut rewritten.rhs).ok_or("right operand is not a selector")?; + if !left.or_matchers.is_empty() || !right.or_matchers.is_empty() { + return Err("selector has an or matcher group"); + } + let constraints = left .matchers .iter() @@ -163,8 +164,11 @@ fn try_propagate( let mut changed = false; for matcher in constraints { - for target in [&mut left.matchers, &mut right.matchers] { - if !target.contains(&matcher) { + for (target, operand) in [ + (&mut left.matchers, &binary.lhs), + (&mut right.matchers, &binary.rhs), + ] { + if !target.contains(&matcher) && preserves_filter(operand, &matcher.name) { target.push(matcher.clone()); changed = true; } @@ -184,9 +188,7 @@ fn matches_every_value(matcher: &Matcher) -> bool { /// by a grouping label drops exactly the matching output series and leaves the remaining /// aggregated values untouched. /// -/// `topk` and `bottomk` rank across a group and carry the input labels through, so filtering -/// before them changes the candidate set. `count_values` adds an output label that does not exist -/// in its input. +/// Ranking aggregates instead retain input labels. `count_values` synthesizes a label. fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { matches!( aggregate.op.id(), @@ -202,22 +204,91 @@ fn partitions_by_grouping_labels(aggregate: &AggregateExpr) -> bool { ) } -/// Returns the matchers of the vector selector an operand scans, together with whether the -/// operand is restricted to the grouping labels an aggregation on its path partitions by, or -/// `None` when the operand's output labels are not proven to be that selector's labels. -fn targeted_selector(expr: &mut Expr) -> Option<(&mut Matchers, bool)> { +/// Ranking operators supported by the planner; limitk and limit_ratio are not implemented. +fn ranks_by_grouping_labels(aggregate: &AggregateExpr) -> bool { + matches!(aggregate.op.id(), token::T_TOPK | token::T_BOTTOMK) +} + +/// Whether a partitioning aggregate proves the operand emits at most one series per combination +/// of its output labels; ranking narrows its input, so it inherits the proof. Uniqueness per +/// *match signature* additionally requires every output label to be a matching label, which the +/// caller checks. +fn has_unique_aggregate_output(expr: &Expr) -> bool { match expr { - Expr::VectorSelector(selector) => Some((&mut selector.matchers, false)), - Expr::Paren(paren) => targeted_selector(&mut paren.expr), - Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { - let (matchers, _) = targeted_selector(&mut aggregate.expr)?; - Some((matchers, aggregate.modifier.is_some())) + Expr::Paren(paren) => has_unique_aggregate_output(&paren.expr), + Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => true, + Expr::Aggregate(aggregate) if ranks_by_grouping_labels(aggregate) => { + has_unique_aggregate_output(&aggregate.expr) } + _ => false, + } +} + +fn vector_operand_is_lhs(binary: &BinaryExpr) -> Option { + if !matches!( + binary.op.id(), + token::T_ADD | token::T_SUB | token::T_MUL | token::T_DIV | token::T_MOD | token::T_POW + ) || binary.modifier.is_some() + { + return None; + } + if matches!(binary.lhs.as_ref(), Expr::NumberLiteral(_)) { + Some(false) + } else if matches!(binary.rhs.as_ref(), Expr::NumberLiteral(_)) { + Some(true) + } else { + None + } +} + +/// Whether adding a matcher on `label` to the scan [`selector_matchers`] finds leaves the +/// operand's output equal to the subset of its unfiltered output that satisfies the matcher. +/// A filter may cross an aggregate only when it removes whole groups: in particular, a label +/// retained by topk is not necessarily one of its partitioning labels. +fn preserves_filter(expr: &Expr, label: &str) -> bool { + match expr { + Expr::VectorSelector(_) => true, + Expr::Paren(paren) => preserves_filter(&paren.expr, label), + Expr::Aggregate(aggregate) => { + let partition_label = match &aggregate.modifier { + Some(LabelModifier::Include(labels)) => labels.labels.iter().any(|x| x == label), + Some(LabelModifier::Exclude(labels)) => labels.labels.iter().all(|x| x != label), + None => false, + }; + partition_label && preserves_filter(&aggregate.expr, label) + } + Expr::Binary(binary) => match vector_operand_is_lhs(binary) { + Some(true) => preserves_filter(&binary.lhs, label), + Some(false) => preserves_filter(&binary.rhs, label), + None => false, + }, + // A rollup emits at most one output series per input series and carries its labels. + Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => true, + _ => false, + } +} + +/// Finds the scanned selector through label-preserving operations. Aggregate partitioning +/// is checked separately for each propagated label. +fn selector_matchers(expr: &mut Expr) -> Option<&mut Matchers> { + match expr { + Expr::VectorSelector(selector) => Some(&mut selector.matchers), + Expr::Paren(paren) => selector_matchers(&mut paren.expr), + Expr::Aggregate(aggregate) + if partitions_by_grouping_labels(aggregate) || ranks_by_grouping_labels(aggregate) => + { + selector_matchers(&mut aggregate.expr) + } + Expr::Binary(binary) => match vector_operand_is_lhs(binary) { + Some(true) => selector_matchers(&mut binary.lhs), + Some(false) => selector_matchers(&mut binary.rhs), + None => None, + }, Expr::Call(call) if LABEL_PRESERVING_RANGE_FUNCTIONS.contains(&call.func.name) => { // Every other argument is a scalar, so position does not matter. let matrix = single_matrix_argument(&call.args.args)?; match call.args.args[matrix].as_mut() { - Expr::MatrixSelector(selector) => Some((&mut selector.vs.matchers, false)), + Expr::MatrixSelector(selector) => Some(&mut selector.vs.matchers), _ => None, } } @@ -226,49 +297,6 @@ fn targeted_selector(expr: &mut Expr) -> Option<(&mut Matchers, bool)> { } } -/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the -/// scanned selector groups by (`by(...)`, the `LabelModifier::Include` of an aggregation). -fn modifier_label_names(expr: &Expr) -> Vec<&String> { - let mut grouped = Vec::new(); - collect_modifier_label_names(expr, true, &mut grouped); - grouped -} - -/// Collects the labels every partitioning aggregation on [`targeted_selector`]' path to the -/// scanned selector excludes from its grouping labels (`without(...)`, the -/// `LabelModifier::Exclude` of an aggregation). -fn modifier_excluded_label_names(expr: &Expr) -> Vec<&String> { - let mut ignored = Vec::new(); - collect_modifier_label_names(expr, false, &mut ignored); - ignored -} - -fn collect_modifier_label_names<'a>(expr: &'a Expr, included: bool, out: &mut Vec<&'a String>) { - match expr { - Expr::Paren(paren) => collect_modifier_label_names(&paren.expr, included, out), - Expr::Aggregate(aggregate) if partitions_by_grouping_labels(aggregate) => { - match aggregate.modifier.as_ref() { - Some(LabelModifier::Include(labels)) if included => out.extend(&labels.labels), - Some(LabelModifier::Exclude(labels)) if !included => out.extend(&labels.labels), - _ => (), - } - collect_modifier_label_names(&aggregate.expr, included, out); - } - _ => (), - } -} - -/// Whether an operand an aggregation partitions by `grouped` labels (with `without(...)`, by -/// every label but `ignored`) still carries `name` as an output label, so that a matcher on -/// `name` reaches it on both sides of that aggregation. -fn crosses_aggregation(name: &String, grouped: &[&String], ignored: &[&String]) -> bool { - if grouped.is_empty() { - !ignored.contains(&name) - } else { - grouped.contains(&name) - } -} - fn single_matrix_argument(args: &[Box]) -> Option { let mut found = None; for (index, arg) in args.iter().enumerate() { @@ -322,7 +350,7 @@ mod tests { #[track_caller] fn assert_rewrite(query: &str, expected: &str) { assert_eq!(rewrite(query).unwrap(), parse(expected).unwrap(), "{query}"); - // The rewrite adds every constraint to both sides, so re-running it is a no-op. + // Re-running the rewrite must be a no-op, including one-way propagation. assert!(rewrite(expected).is_none(), "{expected}"); } @@ -490,12 +518,88 @@ mod tests { ); } + #[test] + fn propagates_through_scalar_arithmetic_and_ranking() { + assert_rewrite( + r#"topk(1, a{host="x"}) / on(host) b"#, + r#"topk(1, a{host="x"}) / on(host) b{host="x"}"#, + ); + assert_rewrite( + r#"a / on(host) bottomk(1, b{host="x"})"#, + r#"a{host="x"} / on(host) bottomk(1, b{host="x"})"#, + ); + assert_rewrite( + r#"(8 * rate(a{host="x"}[5m])) / on(host) topk by(host)(1, b)"#, + r#"(8 * rate(a{host="x"}[5m])) / on(host) topk by(host)(1, b{host="x"})"#, + ); + assert_rewrite( + r#"bottomk without(zone)(2, a / 8) / on(host) b{host="x"}"#, + r#"bottomk without(zone)(2, a{host="x"} / 8) / on(host) b{host="x"}"#, + ); + } + + #[test] + fn propagates_grouped_matching_only_with_a_unique_one_side() { + let cases = [ + ( + r#"(8 * rate(a{host="x"}[5m])) / on(host) group_left topk by(host)(1, max by(host)(b))"#, + r#"(8 * rate(a{host="x"}[5m])) / on(host) group_left topk by(host)(1, max by(host)(b{host="x"}))"#, + vec!["host", "zone"], + vec!["host"], + ), + ( + r#"sum by(host)(a) / on(host) group_right b{host!="x"}"#, + r#"sum by(host)(a{host!="x"}) / on(host) group_right b{host!="x"}"#, + vec!["host"], + vec!["host", "zone"], + ), + ( + r#"a{host=""} / on(host) group_left max by(host)(b)"#, + r#"a{host=""} / on(host) group_left max by(host)(b{host=""})"#, + vec!["host", "zone"], + vec!["host"], + ), + ]; + for (query, expected, left_tags, right_tags) in cases { + assert_eq!( + rewrite_with(query, &left_tags, &right_tags).unwrap(), + parse(expected).unwrap(), + "{query}" + ); + assert!(rewrite_with(expected, &left_tags, &right_tags).is_none()); + } + } + + #[test] + fn preserves_ranking_candidates_and_unproven_cardinality() { + for query in [ + r#"a{host="x"} / on(host) group_left max by(host,zone)(b)"#, + r#"a{host="x"} / on(host) group_left topk by(host)(1, b)"#, + r#"a{host="x"} / on(host) group_left topk(1, max by(host)(b))"#, + r#"topk by(zone)(1, a) / on(host) b{host="x"}"#, + r#"topk by(host)(1, topk(1, a)) / on(host) b{host="x"}"#, + r#"topk by(host)(1, sum by(zone)(a)) / on(host) b{host="x"}"#, + r#"(a + vector(8)) / on(host) b{host="x"}"#, + ] { + assert!(rewrite(query).is_none(), "{query}"); + } + // Even with a unique one-side aggregate, global topk must see every host. + assert!( + rewrite_with( + r#"a{host="x"} / on(host) group_left topk(1, max by(host)(b))"#, + &["host", "zone"], + &["host"] + ) + .is_none() + ); + } + #[test] fn leaves_unproven_semantics_unchanged() { for query in [ r#"a or on(host) b{host="x"}"#, r#"a / on(host) group_left b{host="x"}"#, - r#"sum by(host)(a) / on(host) group_right b{host="x"}"#, + r#"sum by(host,zone)(a) / on(host) group_right b{host="x"}"#, r#"label_replace(a,"host","x","zone",".*") / on(host) b{host="x"}"#, r#"a > on(host) b{host="x"}"#, r#"a / on(host) absent_over_time(b{host="x"}[5m])"#, diff --git a/tests/cases/standalone/common/promql/matching_groups.result b/tests/cases/standalone/common/promql/matching_groups.result new file mode 100644 index 00000000000..183ae6d218c --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.result @@ -0,0 +1,181 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); + +Affected Rows: 5 + +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +Affected Rows: 5 + +-- The result labels are the right operand's tag set rather than what the modifier implies +-- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing +-- only by it are indistinguishable. +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); + ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.Float64(8) * greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| x | 1970-01-01T00:00:00 | 24.0 | +| x | 1970-01-01T00:00:00 | 48.0 | ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.sum(matching_groups_left.greptime_value) / matching_groups_right.greptime_value * Float64(2) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 9.0 | +| x | b | 1970-01-01T00:00:00 | 4.5 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:00 | 6.0 | +| y | 1970-01-01T00:00:00 | 0.36 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); + ++------+----+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+----+-------------------------------------------------------------------------------------------------------+ ++------+----+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + ++------+----+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+----+-------------------------------------------------------------------------------------------------------+ ++------+----+-------------------------------------------------------------------------------------------------------+ + +-- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many +-- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no +-- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands +-- alone, so the recorded output is the unchanged baseline. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 12.0 | +| x | a | 1970-01-01T00:00:00 | 6.0 | +| x | b | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 12.0 | +| x | a | 1970-01-01T00:00:00 | 6.0 | +| x | b | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +DROP TABLE matching_groups_left; + +Affected Rows: 0 + +DROP TABLE matching_groups_right; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/matching_groups.sql b/tests/cases/standalone/common/promql/matching_groups.sql new file mode 100644 index 00000000000..97e46da0437 --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.sql @@ -0,0 +1,65 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +-- The result labels are the right operand's tag set rather than what the modifier implies +-- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing +-- only by it are indistinguishable. +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + +-- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many +-- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no +-- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands +-- alone, so the recorded output is the unchanged baseline. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + +DROP TABLE matching_groups_left; +DROP TABLE matching_groups_right; From aa36f74feb5446e0a58237e8483e28187dbef0a8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 14:38:22 +0000 Subject: [PATCH 55/63] fix(ci): build tests-integration lib with meta-srv/mock (#9299) * fix(ci): build tests-integration lib with meta-srv/mock tests-integration's lib code (src/cluster.rs) uses meta_srv::mocks, but the dependency carrying the mock feature sits in [dev-dependencies]. Builds that only touch the lib, such as the apidoc job's cargo doc --workspace, resolve meta-srv without mock and fail with E0432. --all-targets builds unify dev-dependency features, which is why check, clippy and nextest stayed green. Move the mock-enabled meta-srv entry back to [dependencies]. The other testing features moved out in #9072 are not needed by the lib and stay in [dev-dependencies]. Signed-off-by: Dennis Zhuang * test(repartition): split per-case repartition tests test_repartition_metric ran four format/primary-key-encoding cases in a single test function, and test_repartition_mito ran two format cases. Each case builds its own 3-datanode cluster and runs a full repartition plus GC cycle, so on S3 the metric test took 165-178s against the 180s nextest terminate-after. Merge queue runs failed on it at random. Split each case into its own test. Cases were already independent, so they now run in parallel and each stays far inside the timeout, and a failure points at one encoding instead of four. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- tests-integration/Cargo.toml | 3 +- tests-integration/tests/repartition.rs | 52 +++++++++++++++++++++----- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 7a7d0cba72a..a5ac94924e4 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -65,7 +65,7 @@ hyper-util = { workspace = true, features = ["tokio"] } log-query = { workspace = true } loki-proto.workspace = true meta-client.workspace = true -meta-srv.workspace = true +meta-srv = { workspace = true, features = ["mock"] } mito2.workspace = true object-store.workspace = true operator.workspace = true @@ -111,7 +111,6 @@ hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -meta-srv = { workspace = true, features = ["mock"] } mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 5fe84b19f24..1198e6d23e5 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -45,15 +45,25 @@ macro_rules! repartition_tests { $( paste::item! { mod [] { + // Every case below builds its own cluster and runs a full repartition plus + // GC cycle, which takes around a minute on object-store backends. Keep one + // case per test so a case stays well inside the nextest slow timeout and + // cases run in parallel; do not fold several cases back into one test. + #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_mito >]() { + async fn [< test_repartition_mito_flat >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { common_telemetry::init_default_ut_logging(); - // Cover both storage formats for repartition behavior. - // for flat format $crate::repartition::test_repartition_mito(store_type, true).await; - // for primary key format + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_mito_primary_key >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_mito(store_type, false).await; } } @@ -95,19 +105,41 @@ macro_rules! repartition_tests { } #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_metric >]() { + async fn [< test_repartition_metric_flat_sparse >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { use store_api::codec::PrimaryKeyEncoding; common_telemetry::init_default_ut_logging(); - // Exercise format + primary key encoding matrix for metric engine. - // for flat format with sparse primary key encoding $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Sparse).await; - // for flat format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_flat_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Dense).await; - // for primary key format with sparse primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_sparse >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Sparse).await; - // for primary key format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Dense).await; } } From 045441e3ccce8d474e7e34d9a4e6e5b49f16a74d Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:56:26 +0000 Subject: [PATCH 56/63] feat(json2): support altering JSON2 settings (#9029) * feat(sql): support alter syntax for JSON2 columns Signed-off-by: fys * fix(json2): preserve rows on type hint mismatch during compaction * refactor(json2): simplify alter settings handling * fix(json2): preserve coerced values during compaction * chore: remove unnecessary clone * chor: reduce memory allocations * fix: cargo clippy * chore: update greptime-proto to main branch * refactor(datatypes): unify string handling with other JSON type hints * fix: cr --------- Signed-off-by: fys --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/common/grpc-expr/src/alter.rs | 51 ++- src/common/grpc-expr/src/error.rs | 8 + .../meta/src/ddl/alter_table/executor.rs | 1 + .../src/ddl/alter_table/region_request.rs | 1 + src/common/meta/src/ddl/event/table.rs | 1 + src/datatypes/src/extension/json.rs | 55 ++- src/datatypes/src/json.rs | 384 +++++++++++++++--- src/datatypes/src/vectors/json/array.rs | 105 ++--- src/metric-engine/src/data_region.rs | 1 + src/mito2/src/compaction/json2.rs | 13 +- src/mito2/src/sst/parquet/json_align.rs | 78 +++- src/operator/src/expr_helper.rs | 109 +++-- src/sql/src/parsers/alter_parser.rs | 28 +- src/sql/src/statements/alter.rs | 38 +- src/sql/src/statements/create.rs | 39 +- src/store-api/Cargo.toml | 1 + src/store-api/src/metadata.rs | 36 ++ src/store-api/src/region_request.rs | 166 ++++++++ src/table/src/metadata.rs | 166 +++++++- src/table/src/requests.rs | 11 + .../common/types/json/json2_alter.result | 131 +++++- .../common/types/json/json2_alter.sql | 58 ++- .../json2_alter_type_hints_compaction.result | 83 ++++ .../json2_alter_type_hints_compaction.sql | 36 ++ ...er_type_hints_compaction_conversion.result | 78 ++++ ...alter_type_hints_compaction_conversion.sql | 35 ++ ..._compaction_null_type_hint_mismatch.result | 80 ++++ ...on2_compaction_null_type_hint_mismatch.sql | 35 ++ 30 files changed, 1602 insertions(+), 231 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql create mode 100644 tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result create mode 100644 tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql diff --git a/Cargo.lock b/Cargo.lock index 2c7f3333329..ea55f1b297d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6130,7 +6130,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=0c4beffaa2b3968034915867593a717c6ab0c6b4#0c4beffaa2b3968034915867593a717c6ab0c6b4" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=3d232f406fad574e8c630fb1c304fe2ebb013c9c#3d232f406fad574e8c630fb1c304fe2ebb013c9c" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -14457,6 +14457,7 @@ version = "1.3.0-alpha.1" dependencies = [ "api", "aquamarine", + "arrow-schema 59.2.0", "async-stream", "async-trait", "bytes", diff --git a/Cargo.toml b/Cargo.toml index b4c4c9d7f54..7c5c53d91c8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -161,7 +161,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "0c4beffaa2b3968034915867593a717c6ab0c6b4" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "3d232f406fad574e8c630fb1c304fe2ebb013c9c" } hex = "0.4" hostname = "0.4.0" http = "1" diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index d5214ae5ddb..ffa148c7d0b 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -24,21 +24,23 @@ use api::v1::{ SkippingIndexType as PbSkippingIndexType, column_def, }; use common_query::AddColumnLocation; +use datatypes::json::{JsonSettings, JsonTypeHint}; +use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use table::metadata::{TableId, TableMeta}; use table::requests::{ AddColumnRequest, AlterKind, AlterTableRequest, AnnotationFamily, ModifyColumnTypeRequest, - SetDefaultRequest, SetIndexOption, UnsetIndexOption, + SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, UnsetIndexOption, }; use crate::error::{ self, ColumnNotFoundSnafu, InvalidColumnDefSnafu, InvalidIndexOptionSnafu, - InvalidSetFulltextOptionRequestSnafu, InvalidSetSkippingIndexOptionRequestSnafu, - InvalidSetTableOptionRequestSnafu, InvalidUnsetTableOptionRequestSnafu, - MissingAlterIndexOptionSnafu, MissingFieldSnafu, MissingTableMetaSnafu, - MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu, + InvalidJsonSettingsSnafu, InvalidSetFulltextOptionRequestSnafu, + InvalidSetSkippingIndexOptionRequestSnafu, InvalidSetTableOptionRequestSnafu, + InvalidUnsetTableOptionRequestSnafu, MissingAlterIndexOptionSnafu, MissingFieldSnafu, + MissingTableMetaSnafu, MissingTimestampColumnSnafu, Result, UnknownLocationTypeSnafu, }; const LOCATION_TYPE_FIRST: i32 = LocationType::First as i32; @@ -59,6 +61,34 @@ fn annotation_family_of_keys<'a>( }) } +fn json_settings_from_proto(settings: api::v1::JsonSettings) -> Result { + let type_hints = settings + .type_hints + .into_iter() + .map(|hint| { + let data_type = ConcreteDataType::from( + ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension) + .context(error::ColumnDataTypeSnafu)?, + ); + + Ok(JsonTypeHint { + path: hint.path, + data_type, + // Index configuration is not supported yet, so this is temporarily + // hardcoded to false. + inverted_index: false, + }) + }) + .collect::>>()?; + + JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| { + InvalidJsonSettingsSnafu { + err: err.to_string(), + } + .build() + }) +} + /// Returns the annotation family when `kind` is a SET/UNSET whose keys all /// belong to one family — the alters that only rewrite table metadata and skip /// region dispatch. A mixed batch is an error; never interpret it as "not an @@ -189,6 +219,17 @@ pub fn alter_expr_to_request( columns: modify_column_type_requests, } } + Kind::SetJsonSettings(set_json_settings) => { + let settings = set_json_settings + .settings + .context(MissingFieldSnafu { field: "settings" })?; + AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: set_json_settings.column_name, + settings: json_settings_from_proto(settings)?, + }, + } + } Kind::DropColumns(DropColumns { drop_columns }) => AlterKind::DropColumns { names: drop_columns.into_iter().map(|c| c.name).collect(), }, diff --git a/src/common/grpc-expr/src/error.rs b/src/common/grpc-expr/src/error.rs index ccc40efb5cb..6c9302c39f4 100644 --- a/src/common/grpc-expr/src/error.rs +++ b/src/common/grpc-expr/src/error.rs @@ -73,6 +73,13 @@ pub enum Error { source: api::error::Error, }, + #[snafu(display("Invalid JSON settings: {}", err))] + InvalidJsonSettings { + err: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unknown location type: {}", location_type))] UnknownLocationType { location_type: i32, @@ -188,6 +195,7 @@ impl ErrorExt for Error { | Error::MissingTimestampColumn { .. } => StatusCode::InvalidArguments, Error::MissingField { .. } => StatusCode::InvalidArguments, Error::InvalidColumnDef { source, .. } => source.status_code(), + Error::InvalidJsonSettings { .. } => StatusCode::InvalidArguments, Error::UnknownLocationType { .. } => StatusCode::InvalidArguments, Error::UnknownColumnDataType { .. } | Error::InvalidFulltextIndexColumnType { .. } => { diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs index 53bc4b3d05b..419aa9ae12d 100644 --- a/src/common/meta/src/ddl/alter_table/executor.rs +++ b/src/common/meta/src/ddl/alter_table/executor.rs @@ -333,6 +333,7 @@ fn build_new_table_info( } AlterKind::DropColumns { .. } | AlterKind::ModifyColumnTypes { .. } + | AlterKind::SetJsonSettings { .. } | AlterKind::SetTableOptions { .. } | AlterKind::UnsetTableOptions { .. } | AlterKind::SetAnnotations { .. } diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs index 73edd4c3615..b4f7b679422 100644 --- a/src/common/meta/src/ddl/alter_table/region_request.rs +++ b/src/common/meta/src/ddl/alter_table/region_request.rs @@ -90,6 +90,7 @@ fn create_proto_alter_kind( }))) } Kind::ModifyColumnTypes(x) => Ok(Some(alter_request::Kind::ModifyColumnTypes(x.clone()))), + Kind::SetJsonSettings(x) => Ok(Some(alter_request::Kind::SetJsonSettings(x.clone()))), Kind::DropColumns(x) => { let drop_columns = x .drop_columns diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 670f3c7c1dc..c1ab0e5e5a2 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -193,6 +193,7 @@ pub(crate) fn alter_table_kind_name(kind: &AlterTableKind) -> Option<&'static st AlterTableKind::DropColumns(_) => Some("drop_columns"), AlterTableKind::RenameTable(_) => Some("rename_table"), AlterTableKind::ModifyColumnTypes(_) => Some("modify_column_types"), + AlterTableKind::SetJsonSettings(_) => Some("set_json_settings"), AlterTableKind::SetTableOptions(_) => Some("set_table_options"), AlterTableKind::UnsetTableOptions(_) => Some("unset_table_options"), AlterTableKind::SetIndex(_) => Some("set_index"), diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 3aa0a4c9792..2e86d7791c5 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -24,9 +24,10 @@ use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; -use crate::error::InvalidJson2LayoutSnafu; +use crate::error::{InvalidJson2LayoutSnafu, SerializeSnafu}; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; +use crate::schema::Metadata; /// Aligns JSON2 field types with their built arrays while preserving field and schema metadata. pub fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef { @@ -151,7 +152,7 @@ pub struct JsonMetadata { } impl JsonMetadata { - /// Creates metadata for the JSON2 layout version 2. + /// Creates metadata for the latest JSON2 layout (currently V2). pub fn new(json_settings: JsonSettings) -> Self { Self { json_settings, @@ -300,6 +301,27 @@ impl ExtensionType for Json2ExtensionType { } } +/// Returns JSON2 column metadata with updated settings and the latest layout. +/// +/// Existing SSTs retain their own physical layout metadata. +pub fn json2_metadata_with_updated_settings( + current_metadata: &Metadata, + settings: JsonSettings, +) -> crate::error::Result { + let json_metadata = JsonMetadata::new(settings); + + let mut metadata = current_metadata.clone(); + metadata.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + metadata.insert( + EXTENSION_TYPE_METADATA_KEY.to_string(), + serde_json::to_string(&json_metadata).context(SerializeSnafu)?, + ); + Ok(metadata) +} + /// Checks whether this field is either a legacy JSONB or JSON2 extension type. pub fn is_any_json_extension_type>(field: T) -> bool { let name = field.as_ref().extension_type_name(); @@ -442,6 +464,35 @@ mod tests { Ok(()) } + #[test] + fn test_json2_metadata_with_updated_settings_upgrades_layout() -> crate::error::Result<()> { + for (name, json) in [ + (JsonExtensionType::NAME, r#"{"json_settings":{}}"#), + (Json2ExtensionType::NAME, r#"{"json_settings":{}}"#), + ( + Json2ExtensionType::NAME, + r#"{"json_settings":{},"layout_version":2}"#, + ), + ] { + let metadata = HashMap::from([ + (EXTENSION_TYPE_NAME_KEY.to_string(), name.to_string()), + (EXTENSION_TYPE_METADATA_KEY.to_string(), json.to_string()), + ("other".to_string(), "kept".to_string()), + ]); + let settings = JsonSettings::try_new(vec![], Some(10))?; + let updated = json2_metadata_with_updated_settings(&metadata, settings.clone())?; + assert_eq!(Some("kept"), updated.get("other").map(String::as_str)); + assert_eq!( + Some(Json2ExtensionType::NAME), + updated.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str) + ); + let json_metadata: JsonMetadata = + serde_json::from_str(updated.get(EXTENSION_TYPE_METADATA_KEY).unwrap()).unwrap(); + assert_eq!(JsonMetadata::new(settings), json_metadata); + } + Ok(()) + } + #[test] fn test_parse_json2_physical_layout() -> crate::error::Result<()> { let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index a5b4940e048..cd73e202bf3 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -23,6 +23,7 @@ pub mod value; use std::collections::BTreeMap; use std::collections::btree_map::Entry; +use std::sync::Arc; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; @@ -31,6 +32,7 @@ use snafu::ResultExt; use crate::data_type::ConcreteDataType; use crate::error::{self, InvalidJson2SettingsSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; +use crate::prelude::DataType as _; use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::value::{ListValue, StructValue, Value}; @@ -93,6 +95,15 @@ pub struct JsonTypeHint { pub inverted_index: bool, } +/// Specifies how encoding handles a value that does not match its JSON2 type hint. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TypeHintMismatchPolicy { + /// Returns an error for a type hint mismatch. + Reject, + /// Coerces to the hint type, encoding unconvertible values as null. + CoerceOrNull, +} + /// Context for JSON encoding/decoding that tracks the current key path. #[derive(Clone, Debug)] pub struct JsonContext<'a> { @@ -135,6 +146,30 @@ impl JsonSettings { (self.type_hints, self.max_auto_expanded_paths) } + /// Returns whether the settings are equivalent, ignoring type hint order. + pub fn equivalent(&self, other: &Self) -> bool { + let Self { + type_hints, + max_auto_expanded_paths, + } = self; + let Self { + type_hints: other_hints, + max_auto_expanded_paths: other_max_auto_expanded_paths, + } = other; + + if max_auto_expanded_paths != other_max_auto_expanded_paths + || type_hints.len() != other_hints.len() + { + return false; + } + + let mut hints = type_hints.iter().collect::>(); + let mut other_hints = other_hints.iter().collect::>(); + hints.sort_unstable_by(|a, b| a.path.cmp(&b.path)); + other_hints.sort_unstable_by(|a, b| a.path.cmp(&b.path)); + hints == other_hints + } + /// Decode an encoded StructValue back into a serde_json::Value. pub fn decode(&self, value: Value) -> Result { let mut context = JsonContext { @@ -146,6 +181,15 @@ impl JsonSettings { /// Encode a serde_json::Value into a Value::Json using current settings. pub fn encode(&self, json: Json) -> Result { + self.encode_with_type_hint_mismatch_policy(json, TypeHintMismatchPolicy::Reject) + } + + /// Encodes a serde_json::Value using the given type hint mismatch policy. + pub fn encode_with_type_hint_mismatch_policy( + &self, + json: Json, + policy: TypeHintMismatchPolicy, + ) -> Result { if let Json::Object(object) = &json && object.contains_key(JSON2_REMAINDER_FIELD_NAME) { @@ -160,7 +204,7 @@ impl JsonSettings { path: Vec::new(), settings: self, }; - encode_json_with_context(json, &mut context).map(|v| Value::Json(Box::new(v))) + encode_json_with_context(json, &mut context, policy).map(|v| Value::Json(Box::new(v))) } } @@ -272,36 +316,41 @@ fn with_key_context( } /// Main encoding function with key path tracking -fn encode_json_with_context(json: Json, context: &mut JsonContext) -> Result { +fn encode_json_with_context( + json: Json, + context: &mut JsonContext, + policy: TypeHintMismatchPolicy, +) -> Result { if context.path.is_empty() && !matches!(json, Json::Object(_)) { return UnsupportedJsonTypeSnafu.fail(); } match json { - Json::Object(json_object) => encode_json_object_with_context(json_object, context), - Json::Array(json_array) => encode_json_array_with_context(json_array, context), - _ => encode_json_value_with_context(json, context), + Json::Object(json_object) => encode_json_object_with_context(json_object, context, policy), + Json::Array(json_array) => encode_json_array_with_context(json_array, context, policy), + _ => encode_json_value_with_context(json, context, policy), } } fn encode_json_object_with_context<'a>( json_object: Map, context: &mut JsonContext<'a>, + policy: TypeHintMismatchPolicy, ) -> Result { let mut object = BTreeMap::new(); for (key, value) in json_object { let value = with_key_context(context, &key, |context| { if let Some(hint) = context.type_hint() { - encode_json_value_with_hint(value, hint, context) + encode_json_value_with_hint(value, hint, context, policy) } else { - encode_json_value_with_context(value, context) + encode_json_value_with_context(value, context, policy) } })?; object.insert(key, value.into_variant()); } - fill_missing_type_hints(&mut object, context)?; + fill_missing_type_hints(&mut object, context, policy)?; Ok(JsonValue::new(JsonVariant::Object(object))) } @@ -309,13 +358,14 @@ fn encode_json_object_with_context<'a>( fn fill_missing_type_hints( object: &mut BTreeMap, context: &mut JsonContext, + policy: TypeHintMismatchPolicy, ) -> Result<()> { for hint in &context.settings.type_hints { if hint.path.len() > context.path.len() && hint.path.starts_with(&context.path) { let depth = context.path.len(); let key = &hint.path[depth]; with_key_context(context, key, |context| { - insert_missing_type_hint(object, context, hint, depth) + insert_missing_type_hint(object, context, hint, depth, policy) })?; } } @@ -327,6 +377,7 @@ fn insert_missing_type_hint( field_context: &mut JsonContext, hint: &JsonTypeHint, depth: usize, + policy: TypeHintMismatchPolicy, ) -> Result<()> { let key = &hint.path[depth]; let is_leaf = depth + 1 == hint.path.len(); @@ -341,7 +392,11 @@ fn insert_missing_type_hint( match object.entry(key.clone()) { Entry::Occupied(mut entry) => match entry.get_mut() { JsonVariant::Object(child) => { - insert_missing_type_hint(child, field_context, hint, depth + 1) + insert_missing_type_hint(child, field_context, hint, depth + 1, policy) + } + _ if policy == TypeHintMismatchPolicy::CoerceOrNull => { + entry.insert(JsonValue::null().into_variant()); + Ok(()) } _ => error::InvalidJsonSnafu { value: format!( @@ -354,7 +409,7 @@ fn insert_missing_type_hint( }, Entry::Vacant(entry) => { let mut child = BTreeMap::new(); - insert_missing_type_hint(&mut child, field_context, hint, depth + 1)?; + insert_missing_type_hint(&mut child, field_context, hint, depth + 1, policy)?; entry.insert(JsonVariant::Object(child)); Ok(()) } @@ -365,65 +420,152 @@ fn encode_json_value_with_hint( json: Json, hint: &JsonTypeHint, context: &mut JsonContext, + policy: TypeHintMismatchPolicy, ) -> Result { if json.is_null() { return Ok(JsonValue::null()); } - let invalid_type = || { - error::InvalidJsonSnafu { + let json = match (&hint.data_type, json) { + (ConcreteDataType::String(_), Json::String(value)) => { + return Ok(value.into()); + } + ( + ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_), + Json::Number(value), + ) => { + if let Some(value) = value.as_i64() { + return Ok(value.into()); + } + Json::Number(value) + } + ( + ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_), + Json::Number(value), + ) => { + if let Some(value) = value.as_u64() { + return Ok(value.into()); + } + Json::Number(value) + } + (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(value)) => { + if let Some(value) = value.as_f64() { + return Ok(value.into()); + } + Json::Number(value) + } + (ConcreteDataType::Boolean(_), Json::Bool(value)) => { + return Ok(value.into()); + } + (_, json) => json, + }; + + match policy { + TypeHintMismatchPolicy::Reject => error::InvalidJsonSnafu { value: format!( "JSON value at {} does not match JSON2 type hint {}", context.path.join("."), hint.data_type ), } - .fail() - }; + .fail(), + TypeHintMismatchPolicy::CoerceOrNull => { + let value = coerce_json_value_to_type(json, &hint.data_type); + let value = Json::try_from(value).map_err(|error| { + error::InvalidJsonSnafu { + value: error.to_string(), + } + .build() + })?; + Ok(JsonValue::new(value.into())) + } + } +} - match (&hint.data_type, json) { - (ConcreteDataType::String(_), Json::String(v)) => Ok(v.into()), - ( - ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) - | ConcreteDataType::Int64(_), - Json::Number(v), - ) => match v.as_i64() { - Some(v) => Ok(v.into()), - None => invalid_type(), - }, - ( - ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::UInt64(_), - Json::Number(v), - ) => match v.as_u64() { - Some(v) => Ok(v.into()), - None => invalid_type(), - }, - (ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_), Json::Number(v)) => { - match v.as_f64() { - Some(v) => Ok(v.into()), - None => invalid_type(), +/// Coerces a JSON value using the same semantics as JSON2 query projection. +/// Values that cannot be converted become null. +pub(crate) fn coerce_json_value_to_type(value: Json, to_type: &ConcreteDataType) -> Value { + if value.is_null() { + return Value::Null; + } + + if to_type.is_string() { + let value = match value { + Json::String(value) => value, + value => value.to_string(), + }; + return Value::String(value.into()); + } + + if matches!(to_type, ConcreteDataType::Binary(_)) { + return Value::Binary(encode_serde_json_as_jsonb(value).into()); + } + + if let Some(struct_type) = to_type.as_struct() { + let Json::Object(mut object) = value else { + return Value::Null; + }; + let values = struct_type + .fields() + .iter() + .map(|field| { + object + .remove(field.name()) + .map(|value| coerce_json_value_to_type(value, field.data_type())) + .unwrap_or(Value::Null) + }) + .collect::>(); + return Value::Struct(StructValue::new(values, struct_type.clone())); + } + + if let Some(list_type) = to_type.as_list() { + let Json::Array(values) = value else { + return Value::Null; + }; + let item_type = list_type.item_type().clone(); + let values = values + .into_iter() + .map(|value| coerce_json_value_to_type(value, &item_type)) + .collect::>(); + return Value::List(ListValue::new(values, Arc::new(item_type))); + } + + let value = match value { + Json::Bool(value) => Value::Boolean(value), + Json::Number(value) => { + if let Some(value) = value.as_i64() { + Value::Int64(value) + } else if let Some(value) = value.as_u64() { + Value::UInt64(value) + } else if let Some(value) = value.as_f64() { + Value::Float64(value.into()) + } else { + Value::Null } } - (ConcreteDataType::Boolean(_), Json::Bool(v)) => Ok(v.into()), - _ => invalid_type(), - } + Json::String(value) => Value::String(value.into()), + Json::Array(_) | Json::Object(_) | Json::Null => Value::Null, + }; + to_type.try_cast(value).unwrap_or(Value::Null) } fn encode_json_array_with_context<'a>( json_array: Vec, context: &mut JsonContext<'a>, + policy: TypeHintMismatchPolicy, ) -> Result { let json_array_len = json_array.len(); let mut items = Vec::with_capacity(json_array_len); for (index, value) in json_array.into_iter().enumerate() { let item_value = with_key_context(context, &index.to_string(), |context| { - encode_json_value_with_context(value, context) + encode_json_value_with_context(value, context, policy) })?; items.push(item_value); } @@ -458,7 +600,11 @@ fn encode_json_array_with_context<'a>( } /// Helper function to encode a JSON value to a Value and determine its ConcreteDataType with context -fn encode_json_value_with_context(json: Json, context: &mut JsonContext) -> Result { +fn encode_json_value_with_context( + json: Json, + context: &mut JsonContext, + policy: TypeHintMismatchPolicy, +) -> Result { if context.path.len() >= JSON2_MAX_STRUCTURED_DEPTH && matches!(&json, Json::Object(_) | Json::Array(_)) { @@ -487,8 +633,8 @@ fn encode_json_value_with_context(json: Json, context: &mut JsonContext) -> Resu } } Json::String(s) => Ok(s.into()), - Json::Array(arr) => encode_json_array_with_context(arr, context), - Json::Object(obj) => encode_json_object_with_context(obj, context), + Json::Array(arr) => encode_json_array_with_context(arr, context, policy), + Json::Object(obj) => encode_json_object_with_context(obj, context, policy), } } @@ -597,6 +743,50 @@ mod tests { &struct_value.items()[index] } + #[test] + fn test_json_settings_equivalent() -> Result<()> { + let hints = vec![ + JsonTypeHint { + path: vec!["user".to_string(), "name".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }, + JsonTypeHint { + path: vec!["count".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }, + ]; + let settings = JsonSettings::try_new(hints.clone(), Some(10))?; + let mut reversed = hints.clone(); + reversed.reverse(); + let reordered = JsonSettings::try_new(reversed, Some(10))?; + assert_ne!(settings, reordered); + assert!(settings.equivalent(&reordered)); + assert!(reordered.equivalent(&settings)); + assert!(settings.equivalent(&settings)); + assert_eq!(hints, settings.type_hints()); + assert!(JsonSettings::default().equivalent(&JsonSettings::default())); + + for limit in [None, Some(0), Some(11)] { + assert!(!settings.equivalent(&JsonSettings::try_new(hints.clone(), limit)?)); + } + assert!(!settings.equivalent(&JsonSettings::try_new(vec![hints[0].clone()], Some(10))?)); + + let mut changed = hints.clone(); + changed[0].path = vec!["user".to_string(), "id".to_string()]; + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + + let mut changed = hints.clone(); + changed[0].data_type = ConcreteDataType::int64_datatype(); + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + + let mut changed = hints; + changed[0].inverted_index = true; + assert!(!settings.equivalent(&JsonSettings::try_new(changed, Some(10))?)); + Ok(()) + } + #[test] fn test_json_settings_deserializes_legacy_type_hint_constraints() -> std::result::Result<(), Box> { @@ -755,6 +945,104 @@ mod tests { } } + #[test] + fn test_coerce_or_null_on_type_hint_mismatch() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"kind": "invalid", "message": "kept"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!(struct_field_value(&root, "kind"), &Value::Null); + assert_eq!( + struct_field_value(&root, "message"), + &Value::String("kept".into()) + ); + Ok(()) + } + + #[test] + fn test_coerce_or_null_converts_type_hint_values() -> Result<()> { + let settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["to_string".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }, + JsonTypeHint { + path: vec!["to_int".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }, + ], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"to_string": 123, "to_int": "456", "invalid": "value"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!( + struct_field_value(&root, "to_string"), + &Value::String("123".into()) + ); + assert_eq!(struct_field_value(&root, "to_int"), &Value::Int64(456)); + assert_eq!( + struct_field_value(&root, "invalid"), + &Value::String("value".into()) + ); + Ok(()) + } + + #[test] + fn test_coerce_or_null_on_nested_type_hint_structure_mismatch() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["user".to_string(), "age".to_string()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + None, + )?; + + let result = settings + .encode_with_type_hint_mismatch_policy( + json!({"user": "invalid", "message": "kept"}), + TypeHintMismatchPolicy::CoerceOrNull, + )? + .into_json_inner() + .unwrap(); + let Value::Struct(root) = result else { + panic!("Expected Struct value"); + }; + assert_eq!(struct_field_value(&root, "user"), &Value::Null); + assert_eq!( + struct_field_value(&root, "message"), + &Value::String("kept".into()) + ); + Ok(()) + } + #[test] fn test_encode_rejects_reserved_remainder_field() -> Result<()> { let settings = JsonSettings::default(); diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 08a2fd86555..5b6265c7b37 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -26,15 +26,13 @@ use serde_json::Value; use snafu::{OptionExt, ResultExt}; use crate::arrow_array::{binary_array_value, string_array_value}; -use crate::data_type::ConcreteDataType; +use crate::data_type::{ConcreteDataType, DataType as _}; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; -use crate::json::JsonSettings; -use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; -use crate::prelude::{DataType as _, Value as GreptimeValue}; -use crate::value::{ListValue, StructValue}; +use crate::json::value::decode_json_variant; +use crate::json::{JsonSettings, TypeHintMismatchPolicy, coerce_json_value_to_type}; use crate::vectors::MutableVector; use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type}; use crate::vectors::json::variant::variant_to_json_values; @@ -123,6 +121,23 @@ impl JsonArray<'_> { field: &Field, logical_settings: &JsonSettings, target_layout: &JsonSettings, + ) -> Result { + self.rewrite_to_v2_with_type_hint_mismatch_policy( + field, + logical_settings, + target_layout, + TypeHintMismatchPolicy::Reject, + ) + } + + /// Rewrites a JSON2 array to the specified v2 physical layout using the + /// given type hint mismatch policy. + pub fn rewrite_to_v2_with_type_hint_mismatch_policy( + &self, + field: &Field, + logical_settings: &JsonSettings, + target_layout: &JsonSettings, + policy: TypeHintMismatchPolicy, ) -> Result { let is_v2 = json2_remainder_field(field)?.is_some(); if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) { @@ -141,7 +156,8 @@ impl JsonArray<'_> { if value.is_null() { builder.push_null(); } else { - let value = logical_settings.encode(value)?; + let value = + logical_settings.encode_with_type_hint_mismatch_policy(value, policy)?; builder.try_push_value_ref(&value.as_value_ref())?; } } @@ -319,87 +335,12 @@ fn project_json_values(values: Vec, to_type: &DataType) -> Result Result { - if value.is_null() { - return Ok(GreptimeValue::Null); - } - - if to_type.is_string() { - let value = match value { - Value::String(value) => value, - value => value.to_string(), - }; - return Ok(GreptimeValue::String(value.into())); - } - - if matches!(to_type, ConcreteDataType::Binary(_)) { - return Ok(GreptimeValue::Binary( - encode_serde_json_as_jsonb(value).into(), - )); - } - - if let Some(struct_type) = to_type.as_struct() { - let Value::Object(mut object) = value else { - return Ok(GreptimeValue::Null); - }; - let values = struct_type - .fields() - .iter() - .map(|field| { - object - .remove(field.name()) - .map(|value| project_json_value_to_type(value, field.data_type())) - .transpose() - .map(|value| value.unwrap_or(GreptimeValue::Null)) - }) - .collect::>>()?; - return Ok(GreptimeValue::Struct(StructValue::new( - values, - struct_type.clone(), - ))); - } - - if let Some(list_type) = to_type.as_list() { - let Value::Array(values) = value else { - return Ok(GreptimeValue::Null); - }; - let item_type = list_type.item_type().clone(); - let values = values - .into_iter() - .map(|value| project_json_value_to_type(value, &item_type)) - .collect::>>()?; - return Ok(GreptimeValue::List(ListValue::new( - values, - Arc::new(item_type), - ))); - } - - let value = match value { - Value::Bool(value) => GreptimeValue::Boolean(value), - Value::Number(value) => { - if let Some(value) = value.as_i64() { - GreptimeValue::Int64(value) - } else if let Some(value) = value.as_u64() { - GreptimeValue::UInt64(value) - } else if let Some(value) = value.as_f64() { - GreptimeValue::Float64(value.into()) - } else { - GreptimeValue::Null - } - } - Value::String(value) => GreptimeValue::String(value.into()), - Value::Array(_) | Value::Object(_) => GreptimeValue::Null, - Value::Null => GreptimeValue::Null, - }; - Ok(to_type.try_cast(value).unwrap_or(GreptimeValue::Null)) -} - impl<'a> From<&'a ArrayRef> for JsonArray<'a> { fn from(inner: &'a ArrayRef) -> Self { Self { inner } diff --git a/src/metric-engine/src/data_region.rs b/src/metric-engine/src/data_region.rs index 55b9eaaf683..a11b3e57b54 100644 --- a/src/metric-engine/src/data_region.rs +++ b/src/metric-engine/src/data_region.rs @@ -227,6 +227,7 @@ impl DataRegion { | AlterKind::UnsetRegionOptions { keys: _ } | AlterKind::SetIndexes { options: _ } | AlterKind::UnsetIndexes { options: _ } + | AlterKind::SetJsonSettings { .. } | AlterKind::SyncColumns { column_metadatas: _, } => { diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index f381f80722e..cb2a9d34ec8 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -19,7 +19,9 @@ use arrow_schema::extension::ExtensionType; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, JsonMetadata}; -use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint}; +use datatypes::json::{ + JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint, TypeHintMismatchPolicy, +}; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::json::array::JsonArray; @@ -260,7 +262,7 @@ fn select_dynamic_hints( .filter(|(path, stat)| { !stat.is_type_conflicted // TODO(LFC): Instead of "primitive only", consider retaining stable compound types - // that are safe to write to Parquet, as flush does. Or better, unite the two + // that are safe to write to Parquet, as flush does. Or better, unite the two // selection process. && stat.data_type.is_primitive() && !has_ancestor_path(path) @@ -327,7 +329,12 @@ pub(crate) fn rewrite_json2_batch( }; let array = JsonArray::from(array) - .rewrite_to_v2(field, &plan.logical_settings, &plan.target_layout) + .rewrite_to_v2_with_type_hint_mismatch_policy( + field, + &plan.logical_settings, + &plan.target_layout, + TypeHintMismatchPolicy::CoerceOrNull, + ) .context(ConvertValueSnafu)?; debug_assert_eq!( &json2_physical_data_type(&plan.target_layout), diff --git a/src/mito2/src/sst/parquet/json_align.rs b/src/mito2/src/sst/parquet/json_align.rs index 8af88f6ac0f..93ecffa072f 100644 --- a/src/mito2/src/sst/parquet/json_align.rs +++ b/src/mito2/src/sst/parquet/json_align.rs @@ -22,7 +22,7 @@ use datatypes::arrow::array::{ArrayRef, new_null_array}; use datatypes::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::{JsonMetadata, is_json2_extension_type}; -use datatypes::json::JsonSettings; +use datatypes::json::{JsonSettings, TypeHintMismatchPolicy}; use datatypes::vectors::json::array::JsonArray; use datatypes::vectors::json::json2_physical_data_type; use futures::Stream; @@ -251,10 +251,11 @@ fn rewrite_array( settings: &RewriteSettings, ) -> Result { JsonArray::from(source_array) - .rewrite_to_v2( + .rewrite_to_v2_with_type_hint_mismatch_policy( source_field, &settings.logical_settings, &settings.target_layout, + TypeHintMismatchPolicy::CoerceOrNull, ) .context(DataTypeMismatchSnafu) } @@ -330,7 +331,9 @@ mod tests { Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray, }; use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datatypes::extension::json::Json2ExtensionType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonTypeHint; + use datatypes::prelude::ConcreteDataType; use datatypes::types::parse_string_to_jsonb; use futures::{StreamExt, stream}; @@ -783,6 +786,75 @@ mod tests { assert_eq!(int_array([10, 20]).as_ref(), output.column(3).as_ref()); } + #[tokio::test] + async fn test_rewrite_keeps_rows_with_invalid_json2_settings() { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + inverted_index: false, + }], + Some(0), + ) + .unwrap(); + let target_type = json2_physical_data_type(&settings); + let output_schema = schema([ + Field::new("j", target_type.clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))), + ), + Field::new("value", DataType::Int64, true), + ]); + let source = Arc::new(BinaryArray::from_iter([ + Some(parse_string_to_jsonb(r#"{"kind":"valid"}"#).unwrap()), + Some(parse_string_to_jsonb(r#"{"kind":1}"#).unwrap()), + ])) as ArrayRef; + let input = RecordBatch::try_new( + schema([ + Field::new("j", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("value", DataType::Int64, true), + ]), + vec![source, int_array([10, 20])], + ) + .unwrap(); + let columns = HashMap::from([( + "j".to_string(), + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new(settings.clone())) + .unwrap(), + target_layout: settings, + }, + )]); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true, true], + output_schema, + AlignMode::Rewrite { columns }, + ) + .unwrap(); + + let output = aligner.next().await.unwrap().unwrap(); + assert_eq!(2, output.num_rows()); + assert_eq!( + 10, + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + assert_eq!( + 20, + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(1) + ); + } + #[test] fn test_rewrite_rejects_mismatched_output_layout() { let columns = HashMap::from([( diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index b28af78d383..1d06c24ff7f 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -24,9 +24,10 @@ use api::v1::column_def::{options_from_column_schema, try_as_column_schema}; use api::v1::{ AddColumn, AddColumns, AlterDatabaseExpr, AlterTableExpr, Analyzer, ColumnDataType, ColumnDataTypeExtension, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropColumn, - DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, ModifyColumnType, + DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, + JsonSettings as PbJsonSettings, JsonTypeHint as PbJsonTypeHint, ModifyColumnType, ModifyColumnTypes, RenameTable, SemanticType, SetDatabaseOptions, SetDefaults, SetFulltext, - SetIndex, SetIndexes, SetInverted, SetSkipping, SetTableOptions, + SetIndex, SetIndexes, SetInverted, SetJsonSettings, SetSkipping, SetTableOptions, SkippingIndexType as PbSkippingIndexType, TableName, UnsetDatabaseOptions, UnsetFulltext, UnsetIndex, UnsetIndexes, UnsetInverted, UnsetSkipping, UnsetTableOptions, set_index, unset_index, @@ -36,6 +37,7 @@ use common_error::ext::BoxedError; use common_grpc_expr::util::ColumnExpr; use common_time::Timezone; use datafusion::sql::planner::object_name_to_table_reference; +use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, @@ -782,6 +784,29 @@ pub(crate) fn to_repartition_request( }) } +fn json_settings_to_proto(settings: JsonSettings) -> Result { + let (type_hints, max_auto_expanded_paths) = settings.into_parts(); + let type_hints = type_hints + .into_iter() + .map(|hint| { + let (data_type, datatype_extension) = ColumnDataTypeWrapper::try_from(hint.data_type) + .map(|w| w.to_parts()) + .context(ColumnDataTypeSnafu)?; + + Ok(PbJsonTypeHint { + path: hint.path, + data_type: data_type as i32, + datatype_extension, + }) + }) + .collect::>>()?; + + Ok(PbJsonSettings { + type_hints, + max_auto_expanded_paths, + }) +} + /// Converts a SQL alter table statement into a gRPC alter table expression. pub(crate) fn to_alter_table_expr( alter_table: AlterTable, @@ -830,12 +855,15 @@ pub(crate) fn to_alter_table_expr( } => { let target_type = sql_data_type_to_concrete_data_type(&target_type).context(ParseSqlSnafu)?; + + // Currently disallow modify column type to json2. if target_type.is_json2() { return NotSupportedSnafu { feat: "ALTER TABLE MODIFY COLUMN to JSON2 type", } .fail(); } + let (target_type, target_type_extension) = ColumnDataTypeWrapper::try_from(target_type) .map(|w| w.to_parts()) .context(ColumnDataTypeSnafu)?; @@ -853,6 +881,20 @@ pub(crate) fn to_alter_table_expr( }], }) } + AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + } => { + let settings = match json2_options { + Some(options) => options.build_json_settings().context(ParseSqlSnafu)?, + None => datatypes::json::JsonSettings::new_v2(), + }; + + AlterTableKind::SetJsonSettings(SetJsonSettings { + column_name: column_name.value, + settings: Some(json_settings_to_proto(settings)?), + }) + } AlterTableOperation::DropColumn { name } => AlterTableKind::DropColumns(DropColumns { drop_columns: vec![DropColumn { name: name.value.clone(), @@ -1752,31 +1794,48 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; } #[test] - fn test_alter_json2_is_not_supported() { - for sql in [ - "ALTER TABLE monitor MODIFY COLUMN payload JSON2;", - "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING);", - ] { - let stmt = ParserContext::create_with_dialect( - sql, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .unwrap() - .pop() - .unwrap(); + fn test_to_alter_set_json_settings_expr() { + let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING);"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); - let Statement::AlterTable(alter_table) = stmt else { - unreachable!() - }; - let err = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap_err(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let kind = expr.kind.unwrap(); - assert!(matches!(err, crate::error::Error::NotSupported { .. })); - assert_eq!( - "Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type", - err.to_string() - ); - } + let AlterTableKind::SetJsonSettings(modify) = kind else { + unreachable!() + }; + assert_eq!("payload", modify.column_name); + let settings = modify.settings.as_ref().unwrap(); + assert_eq!(Some(100), settings.max_auto_expanded_paths); + assert_eq!(1, settings.type_hints.len()); + assert_eq!(["service"], &settings.type_hints[0].path[..]); + + let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2;"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let kind = expr.kind.unwrap(); + + let AlterTableKind::SetJsonSettings(modify) = kind else { + unreachable!() + }; + let settings = modify.settings.as_ref().unwrap(); + assert_eq!(Some(100), settings.max_auto_expanded_paths); + assert!(settings.type_hints.is_empty()); } #[test] diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index ca38ec0d65e..3e5c27fda73 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -462,19 +462,19 @@ impl ParserContext<'_> { .context(error::SyntaxSnafu)?; self.parse_alter_table_drop_default(column_name) } else { - let (data_type, json2_options) = - if let Some(json2) = parse_json2_type_and_options(&mut self.parser)? { - json2 - } else { - ( - self.parser.parse_data_type().context(error::SyntaxSnafu)?, - None, - ) - }; + if let Some((_, json2_options)) = + parse_json2_type_and_options(&mut self.parser)? + { + return Ok(AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + }); + } + Ok(AlterTableOperation::ModifyColumnType { column_name, - target_type: data_type, - json2_options, + target_type: self.parser.parse_data_type().context(error::SyntaxSnafu)?, + json2_options: None, }) } } @@ -1576,9 +1576,8 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(alter_table) = statements.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { + let AlterTableOperation::SetJsonSettings { column_name, - target_type, json2_options: Some(options), } = alter_table.alter_operation() else { @@ -1586,7 +1585,6 @@ MODIFY COLUMN attrs JSON2 ( }; assert_eq!("attrs", column_name.value); - assert_eq!("JSON2", target_type.to_string()); assert_eq!(Some(2000), options.max_auto_expanded_paths); assert_eq!(4, options.type_hints.len()); assert_eq!(vec!["user", "id"], options.type_hints[1].path); @@ -1610,7 +1608,7 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(empty) = empty.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { json2_options, .. } = empty.alter_operation() + let AlterTableOperation::SetJsonSettings { json2_options, .. } = empty.alter_operation() else { unreachable!() }; diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index 2f98d8737ae..4fc45c2e69b 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -88,6 +88,11 @@ pub enum AlterTableOperation { target_type: DataType, json2_options: Option, }, + /// `MODIFY JSON2 [json2_options]` + SetJsonSettings { + column_name: Ident, + json2_options: Option, + }, /// `SET =
` SetTableOptions { options: Vec, @@ -265,6 +270,16 @@ impl Display for AlterTableOperation { } Ok(()) } + AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + } => { + write!(f, r#"MODIFY COLUMN {column_name} JSON2"#)?; + if let Some(options) = json2_options { + write!(f, "{options}")?; + } + Ok(()) + } AlterTableOperation::SetTableOptions { options } => { let kvs = options .iter() @@ -433,8 +448,10 @@ impl Display for AlterDatabaseOperation { mod tests { use std::assert_matches; + use super::AlterTableOperation; use crate::dialect::GreptimeDbDialect; use crate::parser::{ParseOptions, ParserContext}; + use crate::statements::create::Json2Options; use crate::statements::statement::Statement; #[test] @@ -503,7 +520,7 @@ ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY, ADD COLUMN } let sql = r"alter table monitor modify column load_15 string;"; - let stmts = + let mut stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) .unwrap(); assert_eq!(1, stmts.len()); @@ -523,6 +540,25 @@ ALTER TABLE monitor MODIFY COLUMN load_15 STRING"#, } } + let Statement::AlterTable(alter_table) = &mut stmts[0] else { + unreachable!(); + }; + let AlterTableOperation::ModifyColumnType { json2_options, .. } = + alter_table.alter_operation_mut() + else { + unreachable!(); + }; + *json2_options = Some(Json2Options { + max_auto_expanded_paths: Some(1), + type_hints: vec![], + }); + assert_eq!( + r#"ALTER TABLE monitor MODIFY COLUMN load_15 STRING( + max_auto_expanded_paths = 1 + )"#, + alter_table.to_string() + ); + let sql = r"alter table monitor drop column load_15;"; let stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 46b5b07c9f7..2310531841c 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -157,6 +157,26 @@ impl Display for Json2Options { } } +impl Json2Options { + pub fn build_json_settings(&self) -> Result { + let type_hints = self + .type_hints + .iter() + .map(|hint| { + Ok(datatypes::json::JsonTypeHint { + path: hint.path.clone(), + data_type: json_type_hint_concrete_data_type(&hint.data_type)?, + inverted_index: hint.inverted_index, + }) + }) + .collect::>>()?; + let max_auto_expanded_paths = self + .max_auto_expanded_paths + .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)); + JsonSettings::try_new(type_hints, max_auto_expanded_paths).map_err(Into::into) + } +} + #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub struct JsonTypeHint { pub path: Vec, @@ -351,24 +371,7 @@ impl ColumnExtensions { return Ok(None); }; - let type_hints = options - .type_hints - .iter() - .map(|hint| { - Ok(datatypes::json::JsonTypeHint { - path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - inverted_index: hint.inverted_index, - }) - }) - .collect::>>()?; - let settings = JsonSettings::try_new( - type_hints, - options - .max_auto_expanded_paths - .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)), - )?; - Ok(Some(settings)) + options.build_json_settings().map(Some) } pub fn set_json_settings(&mut self, settings: JsonSettings) -> Result<()> { diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml index 61ce914fe13..2e5a203912e 100644 --- a/src/store-api/Cargo.toml +++ b/src/store-api/Cargo.toml @@ -10,6 +10,7 @@ workspace = true [dependencies] api.workspace = true aquamarine.workspace = true +arrow-schema.workspace = true async-trait.workspace = true bytes.workspace = true common-base.workspace = true diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 0c663bccc05..21fc3295c44 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -30,6 +30,7 @@ use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use datatypes::arrow; use datatypes::arrow::datatypes::FieldRef; +use datatypes::extension::json::json2_metadata_with_updated_settings; use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef, VectorIndexOptions}; use datatypes::types::TimestampType; use itertools::Itertools; @@ -663,6 +664,10 @@ impl RegionMetadataBuilder { AlterKind::AddColumns { columns } => self.add_columns(columns)?, AlterKind::DropColumns { names } => self.drop_columns(&names), AlterKind::ModifyColumnTypes { columns } => self.modify_column_types(columns)?, + AlterKind::SetJsonSettings { + column_name, + settings, + } => self.set_json_settings(column_name, settings)?, AlterKind::SetIndexes { options } => self.set_indexes(options)?, AlterKind::UnsetIndexes { options } => self.unset_indexes(options)?, AlterKind::SetRegionOptions { options: _ } => { @@ -831,6 +836,37 @@ impl RegionMetadataBuilder { Ok(()) } + fn set_json_settings( + &mut self, + col_name: String, + settings: datatypes::json::JsonSettings, + ) -> Result<()> { + let Some(col_meta) = self + .column_metadatas + .iter_mut() + .find(|col| col.column_schema.name == col_name) + else { + return InvalidRegionRequestSnafu { + region_id: self.region_id, + err: format!("column {col_name} not found"), + } + .fail(); + }; + + let old_metadata = col_meta.column_schema.metadata(); + let new_metadata = + json2_metadata_with_updated_settings(old_metadata, settings).map_err(|err| { + InvalidRegionRequestSnafu { + region_id: self.region_id, + err: err.to_string(), + } + .build() + })?; + + *col_meta.column_schema.mut_metadata() = new_metadata; + Ok(()) + } + fn set_indexes(&mut self, options: Vec) -> Result<()> { let mut set_index_map: HashMap<_, Vec<_>> = HashMap::new(); for option in &options { diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 826ef7d6bc9..113a2393ccc 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -32,6 +32,7 @@ use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint, }; +use arrow_schema::extension::ExtensionType; pub use common_base::AffectedRows; use common_base::readable_size::ReadableSize; use common_grpc::flight::FlightDecoder; @@ -39,6 +40,8 @@ use common_recordbatch::DfRecordBatch; use common_time::range::TimestampRange; use common_time::{TimeToLive, Timestamp}; use datatypes::error::time_index_not_widening_error; +use datatypes::extension::json::Json2ExtensionType; +use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use num_enum::TryFromPrimitive; @@ -795,6 +798,13 @@ pub enum AlterKind { /// Columns to change. columns: Vec, }, + /// Set JSON2 settings of a region column. + SetJsonSettings { + /// Column name. + column_name: String, + /// Target JSON2 settings. + settings: JsonSettings, + }, /// Set region options. SetRegionOptions { options: Vec }, /// Unset region options. @@ -975,6 +985,9 @@ impl AlterKind { col_to_change.validate(metadata)?; } } + AlterKind::SetJsonSettings { column_name, .. } => { + Self::validate_set_json_settings(column_name, metadata)? + } AlterKind::SetRegionOptions { .. } => {} AlterKind::UnsetRegionOptions { .. } => {} AlterKind::SetIndexes { options } => { @@ -1084,6 +1097,18 @@ impl AlterKind { AlterKind::ModifyColumnTypes { columns } => columns .iter() .any(|col_to_change| col_to_change.need_alter(metadata)), + AlterKind::SetJsonSettings { + column_name, + settings, + } => metadata.column_by_name(column_name).is_some_and(|col| { + col.column_schema + .extension_type::() + .ok() + .flatten() + .is_none_or(|extension| { + !extension.metadata().json_settings().equivalent(settings) + }) + }), AlterKind::SetRegionOptions { .. } => true, AlterKind::UnsetRegionOptions { .. } => true, AlterKind::SetIndexes { options, .. } => options @@ -1160,6 +1185,34 @@ impl AlterKind { Ok(()) } + + fn validate_set_json_settings(col_name: &String, metadata: &RegionMetadata) -> Result<()> { + let region_id = metadata.region_id; + + let col = metadata + .column_by_name(col_name) + .with_context(|| InvalidRegionRequestSnafu { + region_id, + err: format!("column {} not found", col_name), + })?; + + ensure!( + col.semantic_type == SemanticType::Field, + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a field column", col_name), + } + ); + ensure!( + col.column_schema.data_type.is_json2(), + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a JSON2 column", col_name), + } + ); + + Ok(()) + } } impl TryFrom for AlterKind { @@ -1183,6 +1236,15 @@ impl TryFrom for AlterKind { .collect::>(); AlterKind::ModifyColumnTypes { columns } } + alter_request::Kind::SetJsonSettings(x) => { + let settings = x.settings.context(InvalidRawRegionRequestSnafu { + err: "missing settings in SetJsonSettings", + })?; + AlterKind::SetJsonSettings { + column_name: x.column_name, + settings: json_settings_from_proto(settings)?, + } + } alter_request::Kind::DropColumns(x) => { let names = x.drop_columns.into_iter().map(|x| x.name).collect(); AlterKind::DropColumns { names } @@ -1484,6 +1546,36 @@ impl From for ModifyColumnType { } } +fn json_settings_from_proto(settings: v1::JsonSettings) -> Result { + let type_hints = settings + .type_hints + .into_iter() + .map(|hint| { + let wrapper = ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension) + .map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + })?; + let data_type = ConcreteDataType::from(wrapper); + + Ok(JsonTypeHint { + path: hint.path, + data_type, + inverted_index: false, + }) + }) + .collect::>>()?; + + JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + }) +} + /// Region option changes used by ALTER requests. /// /// This type is serialized for request persistence. Keep future changes backward @@ -1953,8 +2045,11 @@ mod tests { use api::v1::region::RegionColumnDef; use api::v1::{ColumnDataType, ColumnDef}; use common_time::range::TimestampRange; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend}; + use datatypes::types::JsonType; use super::*; use crate::metadata::RegionMetadataBuilder; @@ -2364,6 +2459,34 @@ mod tests { }, } ); + + let request = RegionAlterRequest::try_from(AlterRequest { + region_id: 0, + schema_version: 1, + kind: Some(alter_request::Kind::SetJsonSettings(v1::SetJsonSettings { + column_name: "payload".to_string(), + settings: Some(v1::JsonSettings { + type_hints: vec![v1::JsonTypeHint { + path: vec!["service".to_string()], + data_type: ColumnDataType::String as i32, + datatype_extension: None, + }], + max_auto_expanded_paths: Some(10), + }), + })), + }) + .unwrap(); + + let AlterKind::SetJsonSettings { + column_name, + settings, + } = request.kind + else { + unreachable!() + }; + assert_eq!("payload", column_name); + assert_eq!(Some(10), settings.max_auto_expanded_paths()); + assert_eq!(1, settings.type_hints().len()); } #[test] @@ -2446,6 +2569,15 @@ mod tests { builder.build().unwrap() } + fn json2_column_schema(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(std::sync::Arc::new( + JsonMetadata::new(settings), + ))); + column_schema + } + #[test] fn test_add_column_validate() { let metadata = new_metadata(); @@ -2730,6 +2862,40 @@ mod tests { assert!(kind.need_alter(&metadata)); } + #[test] + fn test_validate_set_json_settings() { + let mut metadata = new_metadata(); + let current_schema = json2_column_schema("field_0", JsonSettings::new_v2()); + let target_settings = JsonSettings::try_new(vec![], Some(10)).unwrap(); + metadata + .column_metadatas + .iter_mut() + .find(|column| column.column_schema.name == "field_0") + .unwrap() + .column_schema = current_schema.clone(); + + let kind = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: target_settings.clone(), + }; + kind.validate(&metadata).unwrap(); + assert!(kind.need_alter(&metadata)); + + let no_op = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: JsonSettings::new_v2(), + }; + no_op.validate(&metadata).unwrap(); + assert!(!no_op.need_alter(&metadata)); + + AlterKind::SetJsonSettings { + column_name: "tag_0".to_string(), + settings: target_settings, + } + .validate(&metadata) + .unwrap_err(); + } + #[test] fn test_validate_add_columns() { let kind = AlterKind::AddColumns { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 5916dd7ab49..d91771b9c5f 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -22,6 +22,7 @@ use common_query::AddColumnLocation; use datafusion_expr::TableProviderFilterPushDown; use datatypes::error::time_index_not_widening_error; pub use datatypes::error::{Error as ConvertError, Result as ConvertResult}; +use datatypes::extension::json::json2_metadata_with_updated_settings; use datatypes::schema::{ ColumnSchema, FulltextOptions, Schema, SchemaBuilder, SchemaRef, SkippingIndexOptions, }; @@ -41,9 +42,9 @@ use crate::error::{self, Result}; use crate::requests::{ AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, REPARTITION_PARTITION_NUM_HINT_KEY, - SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, has_stable_string_form, - parse_entity_columns, parse_entity_option_key, validate_and_normalize_annotation, - validate_annotation_keys, + SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, TableOptions, UnsetIndexOption, + has_stable_string_form, parse_entity_columns, parse_entity_option_key, + validate_and_normalize_annotation, validate_annotation_keys, }; use crate::table_reference::TableReference; @@ -332,6 +333,7 @@ impl TableMeta { AlterKind::ModifyColumnTypes { columns } => { self.modify_column_types(table_name, columns) } + AlterKind::SetJsonSettings { request } => self.set_json_settings(table_name, request), // No need to rebuild table meta when renaming tables. AlterKind::RenameTable { .. } => Ok(self.new_meta_builder()), AlterKind::SetTableOptions { options } => self.set_table_options(options), @@ -1187,6 +1189,77 @@ impl TableMeta { Ok(meta_builder) } + fn set_json_settings( + &self, + table_name: &str, + req: &SetJsonSettingsRequest, + ) -> Result { + let table_schema = &self.schema; + let idx = table_schema + .column_index_by_name(&req.column_name) + .with_context(|| error::ColumnNotExistsSnafu { + column_name: &req.column_name, + table_name, + })?; + let col = &table_schema.column_schemas()[idx]; + + ensure!( + !self.primary_key_indices.contains(&idx) && table_schema.timestamp_index() != Some(idx), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "Not allowed to change JSON settings for key or timestamp column '{}'", + col.name + ), + } + ); + ensure!( + col.data_type.is_json2(), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!("column '{}' is not a JSON2 column", col.name), + } + ); + + let target_metadata = + json2_metadata_with_updated_settings(col.metadata(), req.settings.clone()).map_err( + |err| { + error::InvalidAlterRequestSnafu { + table: table_name, + err: err.to_string(), + } + .build() + }, + )?; + + let mut cols = table_schema.column_schemas().to_vec(); + *cols[idx].mut_metadata() = target_metadata; + + let mut builder = SchemaBuilder::try_from_columns(cols) + .with_context(|_| error::SchemaBuildSnafu { + msg: format!("Failed to convert column schemas into schema for table {table_name}"), + })? + .version(table_schema.version() + 1); + + for (k, v) in table_schema.metadata().iter() { + builder = builder.add_metadata(k, v); + } + + let new_schema = builder.build().with_context(|_| error::SchemaBuildSnafu { + msg: format!( + "Table {table_name} cannot change JSON settings for column {}", + req.column_name + ), + })?; + + let mut meta_builder = self.new_meta_builder(); + let _ = meta_builder + .schema(Arc::new(new_schema)) + .primary_key_indices(self.primary_key_indices.clone()); + + Ok(meta_builder) + } + /// Split requests into different groups using column location info. fn split_requests_by_column_location<'a>( &self, @@ -1603,9 +1676,12 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use datatypes::data_type::ConcreteDataType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::schema::{ ColumnSchema, FulltextAnalyzer, FulltextBackend, Schema, SchemaBuilder, }; + use datatypes::types::JsonType; use super::*; use crate::Error; @@ -1700,6 +1776,90 @@ mod tests { builder.build().unwrap() } + #[test] + fn test_set_json_settings() { + let schema = Arc::new( + SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + json2_column_schema_v1("payload", JsonSettings::new_v2()), + ]) + .unwrap() + .version(123) + .build() + .unwrap(), + ); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "payload".to_string(), + settings: JsonSettings::try_new(vec![], Some(10)).unwrap(), + }, + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + + let payload = new_meta.schema.column_schema_by_name("payload").unwrap(); + let json_metadata: JsonMetadata = + serde_json::from_str(payload.metadata().get("ARROW:extension:metadata").unwrap()) + .unwrap(); + assert!(json_metadata.is_version_2()); + assert_eq!( + Some(10), + json_metadata.json_settings().max_auto_expanded_paths() + ); + assert_eq!(124, new_meta.schema.version()); + assert_eq!(meta.primary_key_indices, new_meta.primary_key_indices); + } + + fn json2_column_schema_v1(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new( + JsonMetadata::new_v1(settings), + ))); + column_schema + } + + #[test] + fn test_set_json_settings_rejects_non_json2_column() { + let schema = Arc::new(new_test_schema()); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "col2".to_string(), + settings: JsonSettings::new_v2(), + }, + }; + + let err = match meta.builder_with_alter_kind("my_table", &alter_kind) { + Ok(_) => panic!("expected modifying JSON settings on a non-JSON2 column to fail"), + Err(err) => err, + }; + assert!(err.to_string().contains("is not a JSON2 column")); + } + #[test] fn test_modify_time_index_column_type() { let schema = Arc::new(new_test_schema()); diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 1ca99268d67..2c519b32a5e 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -25,6 +25,7 @@ use common_query::AddColumnLocation; use common_time::TimeToLive; use common_time::range::TimestampRange; use datatypes::data_type::ConcreteDataType; +use datatypes::json::JsonSettings; use datatypes::prelude::VectorRef; use datatypes::schema::{ ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, @@ -371,6 +372,13 @@ pub struct ModifyColumnTypeRequest { pub target_type: ConcreteDataType, } +/// Set JSON2 settings request. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SetJsonSettingsRequest { + pub column_name: String, + pub settings: JsonSettings, +} + /// A family of annotation table options: pure metadata markers that no region /// consumes. Setting or unsetting them only rewrites the table's /// `extra_options`, so the alter skips region dispatch entirely. @@ -661,6 +669,9 @@ pub enum AlterKind { ModifyColumnTypes { columns: Vec, }, + SetJsonSettings { + request: SetJsonSettingsRequest, + }, RenameTable { new_table_name: String, }, diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index c23f9e2533f..ab9e39dc777 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -2,37 +2,110 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); Affected Rows: 0 -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type - -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2 (); - -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 2 ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration INT64, + trace_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + ++-----------------------+ +| alter_flushed_sst_num | ++-----------------------+ +| 1 | ++-----------------------+ + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +Affected Rows: 2 + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( @@ -40,10 +113,38 @@ ALTER TABLE application_logs trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | +| 1970-01-01T00:00:00.006 | after-4 | u2 | | r2 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ DROP TABLE application_logs; diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql index c89556c3485..7fce02b76ee 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -2,35 +2,75 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; - -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2 (); +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration INT64, + trace_id STRING ); +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + DROP TABLE application_logs; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result new file mode 100644 index 00000000000..284b7e6f850 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result @@ -0,0 +1,83 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + ++---------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400') | ++---------------------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql new file mode 100644 index 00000000000..45d859d4ece --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql @@ -0,0 +1,36 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result new file mode 100644 index 00000000000..5ba2cf4700a --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result @@ -0,0 +1,78 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + ++--------------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400') | ++--------------------------------------------------------------------------------------+ +| 0 | ++--------------------------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction_conversion; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql new file mode 100644 index 00000000000..de9f211cfb5 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int INT64 + ); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction_conversion; diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result new file mode 100644 index 00000000000..6e0f833ccd6 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result @@ -0,0 +1,80 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + ++----------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400') | ++----------------------------------------------------------------------------------+ +| 0 | ++----------------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +DROP TABLE json2_compaction_null_type_hint_mismatch; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql new file mode 100644 index 00000000000..7ebff9cceb1 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +DROP TABLE json2_compaction_null_type_hint_mismatch; From a9e2a89b7b61877875e522dc32bc134eaf08ed39 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:12:39 +0000 Subject: [PATCH 57/63] fix: keep compaction pruning, metadata, and index work on compact runtime (#9304) * fix: run compaction pruner tasks on compact runtime Signed-off-by: Lei, HUANG * fix: keep compaction metadata and index work on compact runtime Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/access_layer.rs | 8 +- src/mito2/src/cache.rs | 210 +++++++++++++++++++--------- src/mito2/src/cache/file_cache.rs | 7 +- src/mito2/src/read/pruner.rs | 12 +- src/mito2/src/region/opener.rs | 8 +- src/mito2/src/sst/index.rs | 20 ++- src/mito2/src/sst/parquet/reader.rs | 9 +- 7 files changed, 193 insertions(+), 81 deletions(-) diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 337386aa2ff..685cc0d2fc7 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -18,6 +18,7 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use common_base::readable_size::ReadableSize; +use common_runtime::runtime::RuntimeTrait; use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; @@ -348,6 +349,7 @@ impl AccessLayer { write_opts: &WriteOptions, metrics: &mut Metrics, ) -> Result { + let op_type = request.op_type; let region_id = request.metadata.region_id; let region_metadata = request.metadata.clone(); let cache_manager = request.cache_manager.clone(); @@ -426,6 +428,10 @@ impl AccessLayer { // Put parquet metadata to cache manager. if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() { + let runtime = match op_type { + OperationType::Compact => common_runtime::compact_runtime(), + OperationType::Flush => common_runtime::global_runtime(), + }; for sst in &sst_info { if let Some(parquet_metadata) = &sst.file_metadata { let file_id = RegionFileId::new(region_id, sst.file_id); @@ -445,7 +451,7 @@ impl AccessLayer { // Compact cache preparation is best-effort. Run the entire operation in one // detached blocking task so it neither blocks an async worker nor delays the // SST write. - common_runtime::spawn_blocking_global(move || { + runtime.spawn_blocking(move || { match prepare_sst_meta_sync( &file_path, Arc::unwrap_or_clone(parquet_metadata), diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 0d43d31f434..ef150dfb070 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -31,6 +31,8 @@ use std::sync::{Arc, RwLock, Weak}; use bytes::Bytes; use common_base::readable_size::ReadableSize; use common_datasource::compression::CompressionType; +use common_runtime::Runtime; +use common_runtime::runtime::RuntimeTrait; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; @@ -238,7 +240,7 @@ impl SstMetaPreparation { } impl CompactSstMeta { - async fn decode(&self) -> Result> { + async fn decode(&self, runtime: &Runtime) -> Result> { let mut decoded_guard = self.decoded.lock().await; if let Some(decoded) = decoded_guard.upgrade() { return Ok(decoded); @@ -248,33 +250,34 @@ impl CompactSstMeta { let decoded_size = self.decoded_size; let region_metadata = self.region_metadata.upgrade(); let page_index_policy = self.page_index_policy; - let decoded = common_runtime::spawn_blocking_global(move || { - let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( - DecompressObjectSnafu { - compress_type: CompressionType::Zstd, + let decoded = runtime + .spawn_blocking(move || { + let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( + DecompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: "cached SST metadata", + }, + )?; + let bytes = Bytes::from(bytes); + let mut reader = ParquetMetaDataReader::new() + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(page_index_policy); + reader.try_parse(&bytes).context(ReadParquetSnafu { path: "cached SST metadata", - }, - )?; - let bytes = Bytes::from(bytes); - let mut reader = ParquetMetaDataReader::new() - .with_column_index_policy(PageIndexPolicy::Skip) - .with_offset_index_policy(page_index_policy); - reader.try_parse(&bytes).context(ReadParquetSnafu { - path: "cached SST metadata", - })?; - let metadata = reader.finish().context(ReadParquetSnafu { - path: "cached SST metadata", - })?; - CachedSstMeta::try_new_with_page_index_policy( - "cached SST metadata", - metadata, - region_metadata, - page_index_policy, - ) - .map(Arc::new) - }) - .await - .context(JoinSnafu)??; + })?; + let metadata = reader.finish().context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + CachedSstMeta::try_new_with_page_index_policy( + "cached SST metadata", + metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)??; *decoded_guard = Arc::downgrade(&decoded); Ok(decoded) @@ -285,46 +288,50 @@ impl CompactSstMeta { } } -/// Decodes SST metadata without preparing a compact cache entry. +/// Decodes SST metadata on the given blocking runtime without preparing a compact cache entry. pub(crate) async fn decode_sst_meta( file_path: &str, parquet_metadata: ParquetMetaData, region_metadata: Option, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Result> { let file_path = file_path.to_string(); - common_runtime::spawn_blocking_global(move || { - let parquet_metadata = strip_column_indexes(parquet_metadata); - CachedSstMeta::try_new_with_page_index_policy( - &file_path, - parquet_metadata, - region_metadata, - page_index_policy, - ) - .map(Arc::new) - }) - .await - .context(JoinSnafu)? + runtime + .spawn_blocking(move || { + let parquet_metadata = strip_column_indexes(parquet_metadata); + CachedSstMeta::try_new_with_page_index_policy( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)? } -/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime. +/// Decodes SST metadata and attempts to encode both cache representations on the given blocking runtime. pub(crate) async fn prepare_sst_meta( file_path: &str, parquet_metadata: ParquetMetaData, region_metadata: Option, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Result { let file_path = file_path.to_string(); - common_runtime::spawn_blocking_global(move || { - prepare_sst_meta_sync( - &file_path, - parquet_metadata, - region_metadata, - page_index_policy, - ) - }) - .await - .context(JoinSnafu)? + runtime + .spawn_blocking(move || { + prepare_sst_meta_sync( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + }) + .await + .context(JoinSnafu)? } /// Synchronously prepares SST metadata. Callers must run this on a blocking runtime. @@ -671,6 +678,16 @@ pub enum CacheStrategy { } impl CacheStrategy { + /// Returns the runtime for CPU-bound SST metadata work for this request. + pub(crate) fn sst_meta_runtime(&self) -> Runtime { + match self { + CacheStrategy::Compaction(_) => common_runtime::compact_runtime(), + CacheStrategy::EnableAll(_) | CacheStrategy::Disabled => { + common_runtime::global_runtime() + } + } + } + /// Returns whether the SST metadata cache is enabled for this strategy. pub(crate) fn sst_meta_cache_enabled(&self) -> bool { match self { @@ -691,7 +708,12 @@ impl CacheStrategy { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { cache_manager - .get_sst_meta_data(file_id, metrics, page_index_policy) + .get_sst_meta_data( + file_id, + metrics, + page_index_policy, + &self.sst_meta_runtime(), + ) .await } CacheStrategy::Disabled => { @@ -1057,6 +1079,7 @@ impl CacheManager { file_id: RegionFileId, metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option> { let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); let compact = self @@ -1075,7 +1098,7 @@ impl CacheManager { } CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - match compact.decode().await { + match compact.decode(runtime).await { Ok(decoded) => { self.put_sst_meta_data(file_id, decoded.clone()); return Some(decoded); @@ -1101,7 +1124,7 @@ impl CacheManager { let file_cache = write_cache.file_cache(); if self.sst_meta_cache_enabled() { if let Some(metadata) = file_cache - .get_sst_meta_data(key, metrics, page_index_policy) + .get_sst_meta_data(key, metrics, page_index_policy, runtime) .await { metrics.file_cache_hit += 1; @@ -1122,7 +1145,7 @@ impl CacheManager { return Some(decoded); } } else if let Some(decoded) = file_cache - .get_decoded_sst_meta_data(key, metrics, page_index_policy) + .get_decoded_sst_meta_data(key, metrics, page_index_policy, runtime) .await { metrics.file_cache_hit += 1; @@ -1209,7 +1232,7 @@ impl CacheManager { let compact = self .get_compact_sst_meta(&key) .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?; - match compact.decode().await { + match compact.decode(&common_runtime::global_runtime()).await { Ok(metadata) => Some(metadata), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); @@ -2144,7 +2167,12 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2205,14 +2233,24 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); let (metadata, region_metadata) = sst_parquet_meta(); cache.put_parquet_meta_data(file_id, metadata, None); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime(), + ) .await .unwrap(); assert_eq!(region_metadata, cached.region_metadata()); @@ -2230,7 +2268,12 @@ mod tests { cache.remove_parquet_meta_data(file_id); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2247,7 +2290,12 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data( + file_id, + &mut metrics, + Default::default(), + &common_runtime::global_runtime(), + ) .await .unwrap(); assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); @@ -2285,9 +2333,15 @@ mod tests { .set_offset_index(Some(offset_indexes)) .build(); let expected_rows = metadata.file_metadata().num_rows(); - let prepared = prepare_sst_meta("test.parquet", metadata, None, PageIndexPolicy::Required) - .await - .unwrap(); + let prepared = prepare_sst_meta( + "test.parquet", + metadata, + None, + PageIndexPolicy::Required, + &common_runtime::global_runtime(), + ) + .await + .unwrap(); let SstMetaPreparation::Prepared(prepared) = prepared else { panic!("valid metadata should produce a compact cache entry"); }; @@ -2296,7 +2350,12 @@ mod tests { cache.put_prepared_sst_meta(file_id, prepared, false); let mut metrics = MetadataCacheMetrics::default(); let cached = cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Required) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Required, + &common_runtime::global_runtime(), + ) .await .unwrap(); @@ -2377,7 +2436,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime() + ) .await .is_none() ); @@ -2397,7 +2461,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime() + ) .await .is_some() ); @@ -2406,7 +2475,12 @@ mod tests { let mut metrics = MetadataCacheMetrics::default(); assert!( cache - .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Skip) + .get_sst_meta_data( + file_id, + &mut metrics, + PageIndexPolicy::Skip, + &common_runtime::global_runtime() + ) .await .is_some() ); diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 0594a1ac10d..4cb60ed8d48 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -21,6 +21,7 @@ use std::time::{Duration, Instant}; use bytes::Bytes; use common_base::readable_size::ReadableSize; +use common_runtime::Runtime; use common_telemetry::{debug, error, info, warn}; use futures::{AsyncWriteExt, FutureExt, TryStreamExt}; use moka::future::Cache; @@ -627,12 +628,13 @@ impl FileCache { key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option { let file_path = self.inner.cache_file_path(key); let metadata = self .get_parquet_meta_data(key, cache_metrics, page_index_policy) .await?; - match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await { + match prepare_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await { Ok(metadata) => Some(metadata), Err(err) => { CACHE_MISS @@ -653,12 +655,13 @@ impl FileCache { key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + runtime: &Runtime, ) -> Option> { let file_path = self.inner.cache_file_path(key); let metadata = self .get_parquet_meta_data(key, cache_metrics, page_index_policy) .await?; - match decode_sst_meta(&file_path, metadata, None, page_index_policy).await { + match decode_sst_meta(&file_path, metadata, None, page_index_policy, runtime).await { Ok(metadata) => Some(metadata), Err(err) => { CACHE_MISS diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index 19c487ed0b8..ef09f5f48ab 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -350,12 +350,14 @@ impl Pruner { enable_predicate_prefilter, }); - // Spawn worker tasks with their receivers + // Keep pruning and prefetching on the runtime of the originating workload. for (worker_id, rx) in receivers.into_iter().enumerate() { - let inner_clone = inner.clone(); - common_runtime::spawn_query(async move { - Self::worker_loop(worker_id, rx, inner_clone).await; - }); + let worker = Self::worker_loop(worker_id, rx, inner.clone()); + if inner.stream_ctx.input.compaction { + common_runtime::spawn_compact(worker); + } else { + common_runtime::spawn_query(worker); + } } Self { diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index fb5d260fe6f..84a0e13833f 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1247,7 +1247,12 @@ async fn preload_parquet_meta_cache_for_files( let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); if let Some(metadata) = write_cache .file_cache() - .get_sst_meta_data(key, &mut cache_metrics, PageIndexPolicy::Optional) + .get_sst_meta_data( + key, + &mut cache_metrics, + PageIndexPolicy::Optional, + &common_runtime::global_runtime(), + ) .await { let decoded = metadata.decoded(); @@ -1297,6 +1302,7 @@ async fn preload_parquet_meta_cache_for_files( // instead of substituting the region's current schema. None, PageIndexPolicy::Optional, + &common_runtime::global_runtime(), ) .await { diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 0c3ede4fb06..ed3ec36a473 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -57,8 +57,8 @@ use crate::cache::{CacheManagerRef, CacheStrategy}; use crate::config::VectorIndexConfig; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; use crate::error::{ - BuildIndexAsyncSnafu, Error, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, - Result, + BuildIndexAsyncSnafu, Error, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, + RegionTruncatedSnafu, Result, }; use crate::metrics::{ INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL, @@ -774,7 +774,21 @@ impl IndexBuildTask { self.source.file_meta.file_id, )) .await; - match self.index_build(version_control).await { + let result = if self.reason == IndexBuildType::Compact { + let mut task = self.clone(); + // Keep the scheduler slot occupied until the compact runtime finishes the build. + match common_runtime::spawn_compact( + async move { task.index_build(version_control).await }, + ) + .await + { + Ok(result) => result, + Err(err) => Err(err).context(JoinSnafu), + } + } else { + self.index_build(version_control).await + }; + match result { Ok(outcome) => self.on_success(outcome).await, Err(e) => { warn!( diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index f54812f97d4..9584ff0a6af 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -851,7 +851,14 @@ impl ParquetReaderBuilder { let metadata = metadata_loader.load(cache_metrics).await?; let decoded = if self.cache_strategy.sst_meta_cache_enabled() { - let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?; + let metadata = prepare_sst_meta( + file_path, + metadata, + None, + page_index_policy, + &self.cache_strategy.sst_meta_runtime(), + ) + .await?; let decoded = metadata.decoded(); match metadata { SstMetaPreparation::Prepared(metadata) => { From d0f8f4b80c25c3f379acf52bd51399313372ba20 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:12:24 +0000 Subject: [PATCH 58/63] feat: add AI matching, classification, and scoring functions (#9300) * feat: return matching scores from jev Replace the experimental three-argument Boolean function with jev(text, prompt) returning a Float64 probability in [0, 1]. Move threshold comparisons into SQL and update tests and migration examples. Signed-off-by: Lei, HUANG * feat: add Jev choice and score functions Share asynchronous execution across Noul, Choice, and Score. Validate JSON criteria before requests and return typed scalar answers. Add SQL and HTTP mock coverage with usage examples. Signed-off-by: Lei, HUANG * refactor: use generic AI SQL function names Expose ai_match, ai_choose, and ai_score and move their implementation, tests, and usage guide under generic AI names. Document the current unreleased interface without migration history. Signed-off-by: Lei, HUANG * fix: share constant AI criteria within each batch Borrow scalar string arguments and lazily parse constant criteria once per batch. Share the parsed allocation across requests while preserving NULL propagation and batch validation before HTTP calls. Signed-off-by: Lei, HUANG * feat: preserve AI score uncertainty in JSONB results Return score, confidence, and probabilities in criteria-level order from one evaluation. Validate the distribution and preserve provider precision. Add JSON extraction, uncertainty, and single-request regressions, and document confidence-aware ranking. Signed-off-by: Lei, HUANG * docs: explain reuse of volatile AI evaluations Document repeated SELECT and WHERE evaluation costs as N + M requests, and show subquery aliases for reusing scalar or structured AI results without additional model calls. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- docs/how-to/ai-functions.md | 247 ++++++ docs/how-to/jev.md | 119 --- src/common/function/src/function_registry.rs | 22 +- src/common/function/src/scalars.rs | 4 +- src/common/function/src/scalars/ai.rs | 454 +++++++++++ src/common/function/src/scalars/ai/tests.rs | 766 ++++++++++++++++++ src/common/function/src/scalars/jev.rs | 209 ----- src/common/function/src/scalars/jev/tests.rs | 250 ------ .../common/function/ai_functions.result | 144 ++++ .../common/function/ai_functions.sql | 65 ++ .../standalone/common/function/jev.result | 66 -- .../cases/standalone/common/function/jev.sql | 32 - 12 files changed, 1692 insertions(+), 686 deletions(-) create mode 100644 docs/how-to/ai-functions.md delete mode 100644 docs/how-to/jev.md create mode 100644 src/common/function/src/scalars/ai.rs create mode 100644 src/common/function/src/scalars/ai/tests.rs delete mode 100644 src/common/function/src/scalars/jev.rs delete mode 100644 src/common/function/src/scalars/jev/tests.rs create mode 100644 tests/cases/standalone/common/function/ai_functions.result create mode 100644 tests/cases/standalone/common/function/ai_functions.sql delete mode 100644 tests/cases/standalone/common/function/jev.result delete mode 100644 tests/cases/standalone/common/function/jev.sql diff --git a/docs/how-to/ai-functions.md b/docs/how-to/ai-functions.md new file mode 100644 index 00000000000..5c33ffff50e --- /dev/null +++ b/docs/how-to/ai-functions.md @@ -0,0 +1,247 @@ +# AI functions (experimental MVP) + +Three asynchronous SQL scalar functions provide natural-language matching, +classification, and rating. The current backend is TypeSafe's Jev model: + +| Function | Mode | Result | +| --- | --- | --- | +| `ai_match(text, prompt)` | Noul | `Float64` probability in **[0, 1]** that the prompt's statement is true | +| `ai_choose(text, prompt, criteria)` | Choice | `String` containing the selected option's name | +| `ai_score(text, prompt, criteria)` | Score | JSONB object containing `score`, `confidence`, and a `probabilities` array | + +Use scalar results or extracted JSON fields in `SELECT`, comparisons in `WHERE`, +and `ORDER BY`. +All arguments are SQL strings. Any SQL NULL argument produces SQL `NULL` +without validating that row's criteria or calling the API. + +## Start GreptimeDB + +The functions are compiled and registered by the **default-enabled Cargo feature +`ai_functions`**. Enable the Jev backend separately in the server process: + +```sh +export GREPTIMEDB_EXPERIMENTAL_JEV=true +export JEV_API_KEY='' +cargo run -p cmd -- standalone start +``` + +If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. +The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to +`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). +The MVP uses environment variables rather than TOML configuration. +The Cargo feature includes all three SQL functions in the build; the runtime environment +variables enable and configure its API calls. Default compilation does not turn +on external API calls. To enable the feature explicitly, use +`--features ai_functions`. + +## Query + +For a new example table, quote `service` and `message` in the DDL because the +current parser treats them as keywords: + +```sql +CREATE TABLE events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); +``` + +### Matching probability (Noul) + +```sql +SELECT occurred_at, service, message +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match( + (message), + 'The event reports that a payment still failed after retries.' + ) >= 0.8 +ORDER BY occurred_at; +``` + +To return the scores and rank matching events: + +```sql +SELECT occurred_at, message, + ai_match(message, 'The event reports that a payment still failed after retries.') AS score +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' +ORDER BY score DESC NULLS LAST; +``` + +Both arguments are strings; larger scores indicate a stronger match. +`(message)` is an ordinary parenthesized string expression. For several columns, +combine them explicitly, for example `concat(service, ': ', message)`. +Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. + +### Classification (Choice) + +`criteria` is a JSON object with **1 to 255 options**. Each key is an option name; +its value is a description (string, object, or array), or JSON `null` if the name +is sufficient. The function returns the selected key, suitable for filtering +or grouping: + +```sql +SELECT occurred_at, message, + ai_choose(message, 'Which team should handle this event?', + '{"billing":"Payments, invoices, refunds","technical":"Bugs and outages","other":null}') AS team +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z'; +``` + +### Rating (Score) + +`criteria` is a JSON array of **2 to 10 ordered level descriptions**, from low +to high. Descriptions can be strings, objects, or arrays. Level numbers start +at zero. The function returns a JSONB object from one model evaluation: + +```json +{ + "score": 1.05, + "confidence": 0.92, + "probabilities": [0.0, 0.95, 0.05] +} +``` + +- `score` is the provider's probability-weighted mean of the level numbers, in + `[0, N − 1]` for N levels. It may fall between levels. +- `confidence` is the provider's confidence in `[0, 1]`, reflecting how concentrated + the distribution is. It is not a guarantee that the answer is correct. +- `probabilities[i]` is the probability of `criteria[i]`, with exactly N entries in + level order. These are individual level probabilities, not cumulative probabilities. + +Different distributions can have the same score. `[0, 1, 0]` and `[0.5, 0, 0.5]` +both have score `1`: the first concentrates on the middle level, while the second +is split between the extremes. Inspect confidence and the distribution before +treating a mean score as a definite severity level. + +For example, compute the rating once, extract numeric fields, and rank only rows +meeting a chosen confidence threshold: + +```sql +SELECT occurred_at, message, + json_get_float(rating, 'score') AS severity, + json_get_float(rating, 'confidence') AS confidence +FROM ( + SELECT occurred_at, message, + ai_score(message, 'How severe is this event?', + '["No impact to functionality","Degraded service with a workaround","Blocking issue with no workaround"]') AS rating + FROM events + WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' +) AS rated +WHERE json_get_float(rating, 'confidence') >= 0.8 +ORDER BY severity DESC NULLS LAST; +``` + +Use `json_to_string(rating)` to display the full object, or +`json_get_float(rating, 'probabilities[2]')` to read the probability of level 2. +The array positions correspond to the supplied criteria, which provide the level +descriptions. + +## Reusing an evaluation + +AI functions are marked volatile to prevent HTTP calls during planning. Identical +calls written separately in `SELECT` and `WHERE` are evaluated separately. If the +filter evaluates N non-null rows and M rows pass to the projection, this can cost +N + M requests, reaching 2N when all rows pass. + +To return a result and filter by it, compute it once in a subquery and reference +its alias in the outer query: + +```sql +SELECT occurred_at, message, score +FROM ( + SELECT occurred_at, message, + ai_match(message, 'The event reports that a payment still failed after retries.') AS score + FROM events + WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' +) AS scored +WHERE score >= 0.8 +ORDER BY score DESC; +``` + +This also applies to `ai_choose` and `ai_score`. Extract multiple JSON fields from +the aliased `ai_score` result, as in the rating example above, rather than calling +the model again for each field. Filter pushdown preserves the volatile projection +instead of duplicating it into the outer predicate. + +## MVP behavior + +- Each non-null row makes one HTTP request per function invocation: the text is + `state`, and the prompt is the question's `instructions`. Choice and Score pass + the parsed JSON as `criteria`. Calls to different functions are separate requests. +- Criteria for all non-null rows in a batch are validated before sending that + batch's requests. Malformed JSON, invalid description types, or invalid option/ + level counts fail the query locally. +- Answers must have the requested question type. Noul probabilities outside + `[0, 1]`, Choice labels not in the criteria, and Score values outside `[0, N − 1]` + fail the query rather than being clamped or replaced with NULL. Score responses + must also contain a numeric confidence in `[0, 1]` and probabilities for every + level. Each probability must be numeric and in `[0, 1]`; their sum must be within + `1e-6` of 1. Provider values are preserved without renormalizing the distribution + or recomputing the score. +- Up to eight requests run concurrently per expression/batch invocation, not per + query or process. Concurrent partitions and queries can exceed eight requests + in total. Each request has a 30-second timeout. Errors (including rate limits + and invalid responses) fail the query. There is no automatic retry or + cross-query cache. +- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the + candidate set, but SQL does not guarantee left-to-right predicate evaluation or + that `LIMIT` bounds the number of API calls. +- Text passed to the function is sent to the configured TypeSafe endpoint. + The API key remains in the server environment, not in SQL. +- This MVP is intended for standalone use. A distributed deployment would need + the environment on every node evaluating the function and separate validation. + +API reference: + +## Before stabilizing + +The experimental MVP still needs the following controls before stabilization: + +- A process-wide concurrency limit shared by AI function invocations, such as a semaphore. +- Bounded retries with exponential backoff for HTTP `429` and `529`, following + TypeSafe's rate-limit guidance. +- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. + +These are follow-up work, not guarantees provided by the current implementation. + +## Validation + +Check both the default build (where all three functions are registered) and an isolated +`common-function` build without default features. The regular AI function tests use a local +HTTP server and need no API key: + +```sh +cargo nextest run -p common-function +cargo nextest run -p common-function --no-default-features +cargo nextest run -p common-function --no-default-features --features ai_functions +``` + +The sqlness runner explicitly includes `ai_functions` when building its test binary: + +```sh +cargo sqlness bare -t ai_functions +``` + +If using `--bins-dir`, provide a binary built with `ai_functions` (included by +default). CI covers the AI-enabled path through the regular unit tests and +sqlness tests; the feature-off checks above can be run locally when needed. + +An opt-in test calls the real service on synthetic payment events: + +```sh +source ~/.zshrc +GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ + -E 'test(test_ai_match_live)' --run-ignored ignored-only +``` diff --git a/docs/how-to/jev.md b/docs/how-to/jev.md deleted file mode 100644 index 2a21f085946..00000000000 --- a/docs/how-to/jev.md +++ /dev/null @@ -1,119 +0,0 @@ -# Natural-language filtering with Jev (experimental MVP) - -`jev(text, statement, threshold)` returns whether Jev judges the statement true -for the text with probability **greater than or equal to** the threshold. -It is an asynchronous SQL scalar function usable in `WHERE` and `SELECT`. - -## Start GreptimeDB - -Jev is compiled and registered by the **default-enabled Cargo feature -`ai_functions`**. Enable API evaluation separately in the server process: - -```sh -export GREPTIMEDB_EXPERIMENTAL_JEV=true -export JEV_API_KEY='' -cargo run -p cmd -- standalone start -``` - -If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. -The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to -`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). -The MVP uses environment variables rather than TOML configuration. -The Cargo feature includes the SQL function in the build; the runtime environment -variables enable and configure its API calls. Default compilation does not turn -on external API calls. To enable the feature explicitly, use -`--features ai_functions`. - -## Query - -For a new example table, quote `service` and `message` in the DDL because the -current parser treats them as keywords: - -```sql -CREATE TABLE events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); -``` - -```sql -SELECT occurred_at, service, message -FROM events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev( - (message), - 'The event reports that a payment still failed after retries.', - 0.8 - ) -ORDER BY occurred_at; -``` - -The first two arguments are strings and the third is a number in `[0, 1]`. -`(message)` is an ordinary parenthesized string expression. For several columns, -combine them explicitly, for example `concat(service, ': ', message)`. -Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. - -## MVP behavior - -- Each non-null row makes one HTTP request: the text is `state`, and the statement - is a `noul` question's `instructions`. The returned `noul` probability is compared - with the threshold locally; this is not the separate Choice/Score confidence. -- Up to eight requests run concurrently per expression/batch invocation, not per - query or process. Concurrent partitions and queries can exceed eight requests - in total. Each request has a 30-second timeout. Errors (including rate limits - and invalid responses) fail the query. There is no automatic retry or - cross-query cache. -- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the - candidate set, but SQL does not guarantee left-to-right predicate evaluation or - that `LIMIT` bounds the number of API calls. -- Text passed to the function is sent to the configured TypeSafe endpoint. - The API key remains in the server environment, not in SQL. -- This MVP is intended for standalone use. A distributed deployment would need - the environment on every node evaluating the function and separate validation. - -API reference: - -## Before stabilizing - -The experimental MVP still needs the following controls before stabilization: - -- A process-wide concurrency limit shared by Jev invocations, such as a semaphore. -- Bounded retries with exponential backoff for HTTP `429` and `529`, following - TypeSafe's rate-limit guidance. -- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. - -These are follow-up work, not guarantees provided by the current implementation. - -## Validation - -Check both the default build (where `jev` is registered) and an isolated -`common-function` build without default features. Jev's regular tests use a local -HTTP server and need no API key: - -```sh -cargo nextest run -p common-function -cargo nextest run -p common-function --no-default-features -cargo nextest run -p common-function --no-default-features --features ai_functions -``` - -The sqlness runner explicitly includes `ai_functions` when building its test binary: - -```sh -cargo sqlness bare -t jev -``` - -If using `--bins-dir`, provide a binary built with `ai_functions` (included by -default). CI covers the AI-enabled path through the regular unit tests and -sqlness tests; the feature-off checks above can be run locally when needed. - -An opt-in test calls the real service on synthetic payment events: - -```sh -source ~/.zshrc -GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ - -E 'test(test_jev_live)' --run-ignored ignored-only -``` diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index 037ca5b175a..6583f7ba607 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -28,14 +28,14 @@ use crate::aggrs::count_hash::CountHash; use crate::aggrs::vector::VectorFunction as VectorAggrFunction; use crate::function::{Function, FunctionRef}; use crate::function_factory::ScalarFunctionFactory; +#[cfg(feature = "ai_functions")] +use crate::scalars::ai; use crate::scalars::anomaly::AnomalyFunction; use crate::scalars::avg_calc::AvgCalcFunction; use crate::scalars::date::DateFunction; use crate::scalars::expression::ExpressionFunction; use crate::scalars::hll_count::HllCalcFunction; use crate::scalars::ip::IpFunctions; -#[cfg(feature = "ai_functions")] -use crate::scalars::jev::JevFunction; use crate::scalars::json::JsonFunction; use crate::scalars::matches::MatchesFunction; use crate::scalars::matches_term::MatchesTermFunction; @@ -229,7 +229,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| MatchesFunction::register(&function_registry); MatchesTermFunction::register(&function_registry); #[cfg(feature = "ai_functions")] - JevFunction::register(&function_registry); + ai::register(&function_registry); // System and administration functions SystemFunction::register(&function_registry); @@ -384,11 +384,17 @@ mod tests { } #[test] - fn test_jev_registration_matches_feature() { - assert_eq!( - FUNCTION_REGISTRY.get_function("jev").is_some(), - cfg!(feature = "ai_functions") - ); + fn test_ai_registration_matches_feature() { + for name in ["ai_match", "ai_choose", "ai_score"] { + assert_eq!( + FUNCTION_REGISTRY.get_function(name).is_some(), + cfg!(feature = "ai_functions"), + "{name}" + ); + } + for name in ["jev", "jev_choice", "jev_score"] { + assert!(FUNCTION_REGISTRY.get_function(name).is_none(), "{name}"); + } } #[test] diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 61fb7ec632c..4bc4c17f5bb 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -12,13 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "ai_functions")] +pub(crate) mod ai; pub mod anomaly; pub(crate) mod date; pub mod expression; #[cfg(feature = "geo")] pub mod geo; -#[cfg(feature = "ai_functions")] -pub(crate) mod jev; pub mod json; pub mod matches; pub mod matches_term; diff --git a/src/common/function/src/scalars/ai.rs b/src/common/function/src/scalars/ai.rs new file mode 100644 index 00000000000..f2e17dffd36 --- /dev/null +++ b/src/common/function/src/scalars/ai.rs @@ -0,0 +1,454 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::hash::Hash; +use std::marker::PhantomData; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{Array, StringViewArray, new_null_array}; +use arrow::datatypes::DataType; +use async_trait::async_trait; +use datafusion_common::cast::as_string_view_array; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err, not_impl_err}; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use futures::{StreamExt, TryStreamExt, stream}; +use reqwest::Client; +use serde::Serialize; +use serde_json::{Map, Value, json}; + +use crate::function_factory::ScalarFunctionFactory; +use crate::function_registry::FunctionRegistry; + +/// Registers the experimental AI matching, classification, and rating functions. +pub(crate) fn register(registry: &FunctionRegistry) { + AiFunction::::register(registry); + AiFunction::::register(registry); + AiFunction::::register(registry); +} + +/// Shared asynchronous execution for AI functions using the Jev backend. +#[derive(PartialEq, Eq, Hash)] +struct AiFunction { + signature: Signature, + question: PhantomData, + enabled: bool, + api_key: Option, + endpoint: String, + model: String, +} + +struct AiRequest<'a, Q: AiQuestion> { + text: &'a str, + prompt: &'a str, + criteria: Arc, +} + +/// Borrows string arguments without expanding constants to the batch size. +enum StringArgument<'a> { + Scalar(Option<&'a str>), + Array(&'a StringViewArray), +} + +impl<'a> StringArgument<'a> { + fn try_new(arg: &'a ColumnarValue) -> Result { + match arg { + ColumnarValue::Scalar(ScalarValue::Utf8View(value)) => { + Ok(Self::Scalar(value.as_deref())) + } + ColumnarValue::Array(array) => Ok(Self::Array(as_string_view_array(array)?)), + _ => exec_err!("AI functions expect Utf8View arguments"), + } + } + + fn value(&self, row: usize) -> Option<&'a str> { + match self { + Self::Scalar(value) => *value, + Self::Array(array) => { + let array = *array; + array.is_valid(row).then(|| array.value(row)) + } + } + } +} + +impl AiFunction { + fn register(registry: &FunctionRegistry) { + registry.register(ScalarFunctionFactory { + name: Q::NAME.to_string(), + factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), + }); + } + + async fn evaluate(&self, client: &Client, request: AiRequest<'_, Q>) -> Result { + let mut question = json!({ "type": Q::TYPE, "instructions": request.prompt }); + if Q::ARG_COUNT == 3 { + question["criteria"] = json!(request.criteria.as_ref()); + } + let response: Value = client + .post(&self.endpoint) + .json(&json!({ + "model": self.model, + "state": request.text, + "questions": { "matches": question } + })) + .send() + .await + .and_then(|response| response.error_for_status()) + .map_err(|e| exec_datafusion_err!("{} request failed: {e}", Q::NAME))? + .json() + .await + .map_err(|e| exec_datafusion_err!("{} response is not valid JSON: {e}", Q::NAME))?; + + let answer = &response["answers"]["matches"]; + if answer["type"].as_str() != Some(Q::TYPE) { + return exec_err!( + "{} response must contain answers.matches with type {}", + Q::NAME, + Q::TYPE + ); + } + Q::parse_answer(answer, request.criteria.as_ref()) + } + + fn prepare_requests<'a>( + &self, + args: &'a [ColumnarValue], + number_rows: usize, + ) -> Result>>> { + if args.len() != Q::ARG_COUNT { + return exec_err!("{} requires {} arguments", Q::NAME, Q::ARG_COUNT); + } + let args = args + .iter() + .map(StringArgument::try_new) + .collect::>>()?; + let scalar_criteria = args[2..] + .iter() + .all(|arg| matches!(arg, StringArgument::Scalar(_))); + let mut shared_criteria: Option> = None; + + // Validate all non-null rows before making any billable requests. + (0..number_rows) + .map(|row| { + let values: Option> = args.iter().map(|arg| arg.value(row)).collect(); + values + .map(|values| { + // Initialize lazily so NULL rows never validate otherwise invalid criteria. + let criteria = match &shared_criteria { + Some(criteria) => Arc::clone(criteria), + None => { + let criteria = Arc::new(Q::parse_criteria(&values[2..])?); + if scalar_criteria { + shared_criteria = Some(Arc::clone(&criteria)); + } + criteria + } + }; + Ok(AiRequest { + text: values[0], + prompt: values[1], + criteria, + }) + }) + .transpose() + }) + .collect() + } + + fn client(&self) -> Result { + if !self.enabled { + return exec_err!( + "{} is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it", + Q::NAME + ); + } + let key = self + .api_key + .as_deref() + .filter(|key| !key.trim().is_empty()) + .ok_or_else(|| { + exec_datafusion_err!("{} requires the JEV_API_KEY environment variable", Q::NAME) + })?; + let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) + .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; + authorization.set_sensitive(true); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::AUTHORIZATION, authorization); + Client::builder() + .default_headers(headers) + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| exec_datafusion_err!("failed to create {} HTTP client: {e}", Q::NAME)) + } +} + +impl Default for AiFunction { + fn default() -> Self { + Self { + // External model evaluations must not be constant-folded during planning. + signature: Signature::exact( + vec![DataType::Utf8View; Q::ARG_COUNT], + Volatility::Volatile, + ), + question: PhantomData, + enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), + api_key: std::env::var("JEV_API_KEY").ok(), + endpoint: std::env::var("JEV_ENDPOINT") + .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), + model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), + } + } +} + +impl fmt::Debug for AiFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Credentials must never appear in plans or diagnostic output. + f.debug_struct("AiFunction") + .field("name", &Q::NAME) + .field("signature", &self.signature) + .field("enabled", &self.enabled) + .field("model", &self.model) + .finish_non_exhaustive() + } +} + +impl ScalarUDFImpl for AiFunction { + fn name(&self) -> &str { + Q::NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(Q::return_type()) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("{} can only be called from async contexts", Q::NAME) + } +} + +#[async_trait] +impl AsyncScalarUDFImpl for AiFunction { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + let rows = self.prepare_requests(&args.args, args.number_rows)?; + + if rows.iter().all(Option::is_none) { + return Ok(ColumnarValue::Array(new_null_array( + &Q::return_type(), + args.number_rows, + ))); + } + + // Reuse connections within the batch; bound in-flight requests and preserve row order. + let client = self.client()?; + // Collect futures to avoid async_trait's higher-ranked lifetime inference issue. + let requests: Vec<_> = rows + .into_iter() + .map(|row| { + let client = &client; + async move { + match row { + Some(request) => self.evaluate(client, request).await, + None => ScalarValue::try_from(&Q::return_type()), + } + } + }) + .collect(); + // This limit is per expression/batch invocation, not per query or process. + // Concurrent partitions and queries can each have their own in-flight requests. + let answers: Vec = stream::iter(requests).buffered(8).try_collect().await?; + Ok(ColumnarValue::Array(ScalarValue::iter_to_array(answers)?)) + } +} + +/// The request criteria and scalar answer contract for an AI question type. +trait AiQuestion: fmt::Debug + Eq + Hash + Send + Sync + 'static { + type Criteria: Serialize + Send + Sync; + + const NAME: &'static str; + const TYPE: &'static str; + const ARG_COUNT: usize; + + fn return_type() -> DataType; + fn parse_criteria(args: &[&str]) -> Result; + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result; +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Noul; + +impl AiQuestion for Noul { + type Criteria = (); + + const NAME: &'static str = "ai_match"; + const TYPE: &'static str = "noul"; + const ARG_COUNT: usize = 2; + + fn return_type() -> DataType { + DataType::Float64 + } + + fn parse_criteria(_args: &[&str]) -> Result<()> { + Ok(()) + } + + fn parse_answer(answer: &Value, _criteria: &()) -> Result { + numeric_answer(Self::NAME, &answer["noul"], "noul", 1.0) + .map(|probability| ScalarValue::Float64(Some(probability))) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Choice; + +impl AiQuestion for Choice { + type Criteria = Map; + + const NAME: &'static str = "ai_choose"; + const TYPE: &'static str = "choice"; + const ARG_COUNT: usize = 3; + + fn return_type() -> DataType { + DataType::Utf8 + } + + fn parse_criteria(args: &[&str]) -> Result { + let criteria: Value = serde_json::from_str(args[0]) + .map_err(|e| exec_datafusion_err!("ai_choose criteria is not valid JSON: {e}"))?; + match criteria { + Value::Object(options) + if (1..=255).contains(&options.len()) + && options.values().all(|v| v.is_null() || is_description(v)) => + { + Ok(options) + } + _ => exec_err!( + "ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null" + ), + } + } + + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result { + let choice = answer["choice"] + .as_str() + .filter(|choice| criteria.contains_key(*choice)) + .ok_or_else(|| { + exec_datafusion_err!("ai_choose response must contain a choice from the criteria") + })?; + Ok(ScalarValue::Utf8(Some(choice.to_string()))) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct Score; + +impl AiQuestion for Score { + type Criteria = Vec; + + const NAME: &'static str = "ai_score"; + const TYPE: &'static str = "score"; + const ARG_COUNT: usize = 3; + + fn return_type() -> DataType { + DataType::BinaryView + } + + fn parse_criteria(args: &[&str]) -> Result { + let criteria: Value = serde_json::from_str(args[0]) + .map_err(|e| exec_datafusion_err!("ai_score criteria is not valid JSON: {e}"))?; + match criteria { + Value::Array(levels) + if (2..=10).contains(&levels.len()) && levels.iter().all(is_description) => + { + Ok(levels) + } + _ => exec_err!( + "ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays" + ), + } + } + + fn parse_answer(answer: &Value, criteria: &Self::Criteria) -> Result { + let score = numeric_answer( + Self::NAME, + &answer["score"], + "score", + (criteria.len() - 1) as f64, + )?; + let confidence = numeric_answer(Self::NAME, &answer["confidence"], "confidence", 1.0)?; + let probabilities = score_probabilities(&answer["probabilities"], criteria.len())?; + let object = jsonb::Object::from([ + ("score".to_string(), jsonb::Value::from(score)), + ("confidence".to_string(), jsonb::Value::from(confidence)), + ( + "probabilities".to_string(), + jsonb::Value::Array(probabilities.into_iter().map(jsonb::Value::from).collect()), + ), + ]); + Ok(ScalarValue::BinaryView(Some( + jsonb::Value::Object(object).to_vec(), + ))) + } +} + +fn score_probabilities(distribution: &Value, level_count: usize) -> Result> { + if distribution + .as_object() + .is_none_or(|probabilities| probabilities.len() != level_count) + { + return exec_err!( + "ai_score response must contain probabilities for all {level_count} levels" + ); + } + let probabilities = (0..level_count) + .map(|level| { + numeric_answer( + Score::NAME, + &distribution[level.to_string()], + "probability", + 1.0, + ) + }) + .collect::>>()?; + // Allow small rounding differences without renormalizing the provider's distribution. + if (probabilities.iter().sum::() - 1.0).abs() > 1e-6 { + return exec_err!("ai_score response probabilities must sum to 1 within 1e-6"); + } + Ok(probabilities) +} + +fn is_description(description: &Value) -> bool { + matches!( + description, + Value::String(_) | Value::Object(_) | Value::Array(_) + ) +} + +fn numeric_answer(name: &str, answer: &Value, field: &str, max: f64) -> Result { + answer + .as_f64() + .filter(|score| (0.0..=max).contains(score)) + .ok_or_else(|| { + exec_datafusion_err!("{name} response must contain a finite {field} in [0, {max}]") + }) +} diff --git a/src/common/function/src/scalars/ai/tests.rs b/src/common/function/src/scalars/ai/tests.rs new file mode 100644 index 00000000000..0b001e22e7a --- /dev/null +++ b/src/common/function/src/scalars/ai/tests.rs @@ -0,0 +1,766 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use axum::http::{HeaderMap, StatusCode}; +use axum::routing::post; +use axum::{Json, Router}; +use datafusion::assert_batches_eq; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +use super::*; +use crate::function::FunctionContext; +use crate::function_registry::FUNCTION_REGISTRY; + +struct MockServer { + endpoint: String, + task: JoinHandle<()>, +} + +impl MockServer { + async fn start(router: Router) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + Self { endpoint, task } + } + + fn function(&self) -> AiFunction { + AiFunction { + enabled: true, + api_key: Some("test-key".to_string()), + endpoint: self.endpoint.clone(), + model: "test-model".to_string(), + ..Default::default() + } + } +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn context(function: AiFunction) -> SessionContext { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); + let json_get_float = FUNCTION_REGISTRY.get_function("json_get_float").unwrap(); + ctx.register_udf(json_get_float.provide(FunctionContext::default())); + ctx +} + +#[test] +fn test_ai_scalar_criteria_share_allocation_and_skip_null_rows() { + fn assert_shared_criteria(criteria: &str) { + let function = AiFunction::::default(); + for (texts, criteria) in [ + (vec![None, Some("first"), Some("second")], criteria), + (vec![None, None], "not json"), + (vec![], "not json"), + ] { + let number_rows = texts.len(); + let expected_requests = texts.iter().flatten().count(); + let args = vec![ + ColumnarValue::Array(Arc::new(StringViewArray::from(texts))), + ColumnarValue::Scalar(ScalarValue::Utf8View(Some("prompt".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8View(Some(criteria.to_string()))), + ]; + let requests = function.prepare_requests(&args, number_rows).unwrap(); + assert_eq!(requests.len(), number_rows); + assert_eq!(requests.iter().flatten().count(), expected_requests); + + // Regression: parsed constant criteria must not be duplicated for each row. + let mut criteria = requests.iter().flatten().map(|request| &request.criteria); + if let Some(first) = criteria.next() { + assert!(criteria.all(|other| Arc::ptr_eq(first, other))); + } + } + } + + assert_shared_criteria::(r#"{"billing":null,"technical":"errors"}"#); + assert_shared_criteria::(r#"["low","high"]"#); +} + +#[tokio::test] +async fn test_ai_match_sql_returns_and_filters_noul_probability() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post( + |headers: HeaderMap, Json(request): Json| async move { + assert_eq!(headers["authorization"], "Bearer test-key"); + assert_eq!(request["model"], "test-model"); + assert_eq!(request["questions"]["matches"]["type"], "noul"); + // The prompt is passed through verbatim. + assert_eq!( + request["questions"]["matches"]["instructions"], + "payment failed" + ); + let probability = match request["state"].as_str().unwrap() { + "failed" => { + // Complete this row after later rows to exercise result alignment. + tokio::time::sleep(Duration::from_millis(20)).await; + 0.95 + } + "boundary" => 0.8, + "recovered" => 0.1, + "impossible" => 0.0, + "certain" => 1.0, + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) + }, + ), + )) + .await; + let ctx = context(server.function::()); + let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL), (5, 'impossible'), (6, 'certain')) AS events(id, message)"; + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE ai_match((message), 'payment failed') >= 0.8 ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 6 |", "+----+" + ], + &batches + ); + + let batches = ctx + .sql(&format!( + "SELECT id, ai_match(message, 'payment failed') AS score FROM {events} ORDER BY score DESC NULLS LAST" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+----+-------+", + "| id | score |", + "+----+-------+", + "| 6 | 1.0 |", + "| 1 | 0.95 |", + "| 2 | 0.8 |", + "| 3 | 0.1 |", + "| 5 | 0.0 |", + "| 4 | |", + "+----+-------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_match_nulls_need_no_api_key() { + let ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let batches = ctx + .sql("SELECT ai_match(NULL, 'condition') AS a, ai_match('text', NULL) AS b") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(0).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+---+---+", + "| a | b |", + "+---+---+", + "| | |", + "+---+---+" + ], + &batches + ); + + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); + + let ctx = context(AiFunction:: { + enabled: false, + api_key: None, + ..Default::default() + }); + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), + "{error}" + ); +} + +#[tokio::test] +async fn test_ai_match_bad_api_responses_fail_the_query() { + for (status, body, expected) in [ + ( + StatusCode::TOO_MANY_REQUESTS, + "{}", + "ai_match request failed", + ), + ( + StatusCode::OK, + "not json", + "ai_match response is not valid JSON", + ), + ( + StatusCode::OK, + r#"{"answers":{}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":-0.1}}}"#, + "ai_match response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, + "ai_match response must contain", + ), + ] { + let server = MockServer::start( + Router::new().route("/v1/systemone", post(move || async move { (status, body) })), + ) + .await; + let ctx = context(server.function::()); + let error = ctx + .sql("SELECT ai_match('text', 'condition')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + } +} + +#[tokio::test] +async fn test_ai_choose_returns_and_filters_supplied_labels() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(|Json(request): Json| async move { + let question = &request["questions"]["matches"]; + assert_eq!(question["type"], "choice"); + assert_eq!(question["instructions"], "Route the ticket"); + let choice = match request["state"].as_str().unwrap() { + "refund" => { + assert_eq!(question["criteria"], json!({"billing": null, "technical": {"examples": ["crash"]}})); + tokio::time::sleep(Duration::from_millis(20)).await; + "billing" + } + "crash" => { + assert_eq!(question["criteria"], json!({"billing": ["refunds"], "technical": "errors"})); + "technical" + } + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "choice", "choice": choice, "confidence": 0.7}}})) + }), + )) + .await; + let ctx = context(server.function::()); + let events = r#"(VALUES + (1, 'refund', 'Route the ticket', '{"billing":null,"technical":{"examples":["crash"]}}'), + (2, 'crash', 'Route the ticket', '{"billing":["refunds"],"technical":"errors"}'), + (3, NULL, 'Route the ticket', 'invalid JSON'), + (4, 'refund', NULL, 'invalid JSON'), + (5, 'refund', 'Route the ticket', NULL) + ) AS events(id, message, prompt, criteria)"#; + let batches = ctx + .sql(&format!( + "SELECT id, ai_choose(message, prompt, criteria) AS team FROM {events} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Utf8); + assert_batches_eq!( + [ + "+----+-----------+", + "| id | team |", + "+----+-----------+", + "| 1 | billing |", + "| 2 | technical |", + "| 3 | |", + "| 4 | |", + "| 5 | |", + "+----+-----------+" + ], + &batches + ); + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE ai_choose(message, prompt, criteria) = 'billing'" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); + + let batches = ctx + .sql( + r#"SELECT id, ai_choose(message, 'Route the ticket', + '{"billing":null,"technical":{"examples":["crash"]}}') AS team + FROM (VALUES (1, 'refund'), (2, NULL), (3, 'refund')) AS events(id, message) + ORDER BY id"#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+---------+", + "| id | team |", + "+----+---------+", + "| 1 | billing |", + "| 2 | |", + "| 3 | billing |", + "+----+---------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_score_returns_fractional_scores_on_each_rows_scale() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(|Json(request): Json| async move { + let question = &request["questions"]["matches"]; + assert_eq!(question["type"], "score"); + assert_eq!(question["instructions"], "Rate severity"); + let (score, probabilities) = match request["state"].as_str().unwrap() { + "low" => { + assert_eq!(question["criteria"], json!(["low", {"description": "medium"}, ["high"]])); + tokio::time::sleep(Duration::from_millis(20)).await; + (0.0, json!({"0":1.0,"1":0.0,"2":0.0})) + } + "fractional" => (1.25, json!({"2":0.25,"0":0.0,"1":0.75})), + "high" => (2.0, json!({"0":0.0,"1":0.0,"2":1.0})), + "binary" => { + assert_eq!(question["criteria"], json!(["low", "high"])); + (1.0, json!({"0":0.0,"1":1.0})) + } + "ten levels" => { + assert_eq!(question["criteria"].as_array().unwrap().len(), 10); + let probabilities: Map = (0..10) + .map(|level| (level.to_string(), json!(if level == 9 { 1.0 } else { 0.0 }))) + .collect(); + (9.0, Value::Object(probabilities)) + } + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "score", "score": score, "confidence": 0.4, "probabilities": probabilities}}})) + }), + )) + .await; + let ctx = context(server.function::()); + let batches = ctx + .sql( + r#" + SELECT id, json_get_float(ai_score(message, 'Rate severity', criteria), 'score') AS severity + FROM (VALUES + (1, 'low', '["low", {"description":"medium"}, ["high"]]'), + (2, 'fractional', '["low", "medium", "high"]'), + (3, 'high', '["low", "medium", "high"]'), + (4, 'binary', '["low", "high"]'), + (5, 'ten levels', '["a","b","c","d","e","f","g","h","i","j"]'), + (6, NULL, 'invalid JSON'), + (7, 'high', NULL) + ) AS events(id, message, criteria) + ORDER BY severity DESC NULLS LAST, id + "#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].schema().field(1).data_type(), &DataType::Float64); + assert_batches_eq!( + [ + "+----+----------+", + "| id | severity |", + "+----+----------+", + "| 5 | 9.0 |", + "| 3 | 2.0 |", + "| 2 | 1.25 |", + "| 4 | 1.0 |", + "| 1 | 0.0 |", + "| 6 | |", + "| 7 | |", + "+----+----------+" + ], + &batches + ); + + let batches = ctx + .sql( + r#"SELECT id, json_get_float(ai_score(message, 'Rate severity', + '["low",{"description":"medium"},["high"]]'), 'score') AS severity + FROM (VALUES (1, 'low'), (2, 'fractional'), (3, 'high'), (4, NULL)) AS events(id, message) + ORDER BY id"#, + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+----------+", + "| id | severity |", + "+----+----------+", + "| 1 | 0.0 |", + "| 2 | 1.25 |", + "| 3 | 2.0 |", + "| 4 | |", + "+----+----------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_ai_score_preserves_uncertainty_in_one_evaluation() { + let calls = Arc::new(AtomicUsize::new(0)); + let request_calls = calls.clone(); + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move |Json(request): Json| { + request_calls.fetch_add(1, Ordering::Relaxed); + async move { + let (confidence, probabilities) = match request["state"].as_str().unwrap() { + "medium" => (1.0, json!({"2":0.0,"0":0.0,"1":1.0})), + "split" => (0.25, json!({"2":0.5,"0":0.5,"1":0.0})), + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers":{"matches":{ + "type":"score", "score":1.0, "confidence":confidence, + "probabilities":probabilities + }}})) + } + }), + )) + .await; + let ctx = context(server.function::()); + let rated = r#"(SELECT id, ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating + FROM (VALUES (1, 'medium'), (2, 'split'), (3, NULL)) AS events(id, message)) AS rated"#; + let batches = ctx + .sql(&format!( + "SELECT id, json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[0]') AS p0, + json_get_float(rating, 'probabilities[1]') AS p1, + json_get_float(rating, 'probabilities[2]') AS p2 + FROM {rated} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+-------+------------+-----+-----+-----+", + "| id | score | confidence | p0 | p1 | p2 |", + "+----+-------+------------+-----+-----+-----+", + "| 1 | 1.0 | 1.0 | 0.0 | 1.0 | 0.0 |", + "| 2 | 1.0 | 0.25 | 0.5 | 0.0 | 0.5 |", + "| 3 | | | | | |", + "+----+-------+------------+-----+-----+-----+" + ], + &batches + ); + assert_eq!(calls.load(Ordering::Relaxed), 2); + + let batches = ctx + .sql(&format!( + "SELECT id FROM {rated} WHERE json_get_float(rating, 'confidence') >= 0.8 + ORDER BY json_get_float(rating, 'score') DESC" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); + assert_eq!(calls.load(Ordering::Relaxed), 4); +} + +#[test] +fn test_ai_score_preserves_provider_precision() { + let answer = json!({ + "score":1.6, "confidence":0.92, + "probabilities":{"2":0.6999995,"0":0.1,"1":0.2} + }); + let criteria = vec![json!("low"), json!("medium"), json!("high")]; + let ScalarValue::BinaryView(Some(bytes)) = Score::parse_answer(&answer, &criteria).unwrap() + else { + panic!("expected a JSONB object"); + }; + let object: Value = + serde_json::from_str(&jsonb::from_slice(&bytes).unwrap().to_string()).unwrap(); + assert_eq!(object.as_object().unwrap().len(), 3); + assert_eq!(object["score"].as_f64(), Some(1.6)); + assert_eq!(object["confidence"].as_f64(), Some(0.92)); + let probabilities: Vec = object["probabilities"] + .as_array() + .unwrap() + .iter() + .map(|probability| probability.as_f64().unwrap()) + .collect(); + assert_eq!(probabilities, vec![0.1, 0.2, 0.6999995]); +} + +#[tokio::test] +async fn test_ai_score_rejects_invalid_uncertainty() { + for (field, value) in [ + ("confidence", Value::Null), + ("confidence", json!(-0.1)), + ("confidence", json!(1.1)), + ("confidence", json!("0.9")), + ("probabilities", Value::Null), + ("probabilities", json!([0.0, 1.0, 0.0])), + ("probabilities", json!({"0":1.0})), + ("probabilities", json!({"0":0.5,"2":0.5,"3":0.0})), + ("probabilities", json!({"0":0.0,"1":1.0,"2":0.0,"3":0.0})), + ("probabilities", json!({"0":-0.1,"1":0.6,"2":0.5})), + ("probabilities", json!({"0":0.0,"1":1.1,"2":0.0})), + ("probabilities", json!({"0":0.0,"1":0.95,"2":"0.05"})), + ("probabilities", json!({"0":0.0,"1":0.0,"2":0.0})), + ("probabilities", json!({"0":0.5,"1":0.5,"2":0.5})), + ] { + let mut answer = json!({ + "type":"score", "score":1.0, "confidence":0.9, + "probabilities":{"0":0.0,"1":1.0,"2":0.0} + }); + answer[field] = value; + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move || { + let answer = answer.clone(); + async move { Json(json!({"answers":{"matches":answer}})) } + }), + )) + .await; + let ctx = context(server.function::()); + let error = ctx + .sql(r#"SELECT ai_score('text', 'prompt', '["low","medium","high"]')"#) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("ai_score response"), "{error}"); + } +} + +#[tokio::test] +async fn test_ai_invalid_criteria_fail_before_requests() { + let choice_ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let score_ctx = context(AiFunction:: { + enabled: true, + api_key: None, + ..Default::default() + }); + let too_many_choices: Map = (0..256) + .map(|index| (index.to_string(), Value::Null)) + .collect(); + for (ctx, name, valid, invalid) in [ + ( + &choice_ctx, + "ai_choose", + r#"{"billing":null}"#, + vec![ + "not json".to_string(), + "null".to_string(), + "{}".to_string(), + "[]".to_string(), + r#"{"billing":1}"#.to_string(), + r#"{"billing":true}"#.to_string(), + Value::Object(too_many_choices).to_string(), + ], + ), + ( + &score_ctx, + "ai_score", + r#"["low","high"]"#, + vec![ + "not json".to_string(), + "null".to_string(), + "{}".to_string(), + "[]".to_string(), + r#"["low"]"#.to_string(), + r#"["low",null]"#.to_string(), + r#"["low",1]"#.to_string(), + r#"["low",true]"#.to_string(), + json!(vec!["level"; 11]).to_string(), + ], + ), + ] { + let error = ctx + .sql(&format!( + "SELECT {name}(message, 'prompt', 'not json') FROM (VALUES (NULL), ('text')) AS input(message)" + )) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error.to_string().contains(&format!("{name} criteria")), + "{error}" + ); + for criteria in invalid { + // A later invalid row must fail before even creating a client for the first row. + let error = ctx.sql(&format!( + "SELECT {name}('text', 'prompt', criteria) FROM (VALUES ('{valid}'), ('{criteria}')) AS input(criteria)" + )).await.unwrap().collect().await.unwrap_err(); + assert!( + error.to_string().contains(&format!("{name} criteria")), + "{error}" + ); + } + } +} + +#[tokio::test] +async fn test_ai_choose_and_score_reject_invalid_answers() { + for (name, criteria, answer) in [ + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"score","choice":"billing"}), + ), + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"choice","choice":"unknown"}), + ), + ( + "ai_choose", + r#"{"billing":null}"#, + json!({"type":"choice","choice":1}), + ), + ("ai_choose", r#"{"billing":null}"#, json!({"type":"choice"})), + ( + "ai_score", + r#"["low","medium","high"]"#, + json!({"type":"noul","score":0.5}), + ), + ( + "ai_score", + r#"["low","medium","high"]"#, + json!({"type":"score","score":2.1}), + ), + ( + "ai_score", + r#"["low","high"]"#, + json!({"type":"score","score":-0.1}), + ), + ( + "ai_score", + r#"["low","high"]"#, + json!({"type":"score","score":"NaN"}), + ), + ("ai_score", r#"["low","high"]"#, json!({"type":"score"})), + ] { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post(move || { + let answer = answer.clone(); + async move { Json(json!({"answers":{"matches":answer}})) } + }), + )) + .await; + let ctx = context(server.function::()); + ctx.register_udf( + AsyncScalarUDF::new(Arc::new(server.function::())).into_scalar_udf(), + ); + let error = ctx + .sql(&format!("SELECT {name}('text', 'prompt', '{criteria}')")) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains(&format!("{name} response must contain")), + "{error}" + ); + } +} + +#[tokio::test] +#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] +async fn test_ai_match_live() { + let ctx = SessionContext::new(); + let function = FUNCTION_REGISTRY.get_function("ai_match").unwrap(); + ctx.register_udf(function.provide(FunctionContext::default())); + let batches = ctx.sql( + "SELECT id FROM (VALUES + (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), + (2, 'Payment succeeded on the second retry; the payment is complete.'), + (3, 'User logged in successfully.') + ) AS events(id, message) + WHERE ai_match(message, 'The event reports that a payment still failed after retries.') >= 0.8 + ORDER BY id", + ).await.unwrap().collect().await.unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); +} diff --git a/src/common/function/src/scalars/jev.rs b/src/common/function/src/scalars/jev.rs deleted file mode 100644 index 1c6475e6fac..00000000000 --- a/src/common/function/src/scalars/jev.rs +++ /dev/null @@ -1,209 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#[cfg(test)] -mod tests; - -use std::fmt; -use std::sync::Arc; -use std::time::Duration; - -use arrow::array::BooleanArray; -use arrow::datatypes::DataType; -use async_trait::async_trait; -use datafusion_common::cast::{as_float64_array, as_string_view_array}; -use datafusion_common::utils::take_function_args; -use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err}; -use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; -use futures::{StreamExt, TryStreamExt, stream}; -use reqwest::Client; -use serde_json::{Value, json}; - -use crate::function_factory::ScalarFunctionFactory; -use crate::function_registry::FunctionRegistry; - -/// Experimental natural-language predicate: `jev(text, statement, threshold)`. -/// Each non-null row is a Noul question; its probability is compared with `>=`. -#[derive(PartialEq, Eq, Hash)] -pub(crate) struct JevFunction { - signature: Signature, - enabled: bool, - api_key: Option, - endpoint: String, - model: String, -} - -impl JevFunction { - pub(crate) fn register(registry: &FunctionRegistry) { - registry.register(ScalarFunctionFactory { - name: "jev".to_string(), - factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), - }); - } - - async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result { - let response: Value = client - .post(&self.endpoint) - .json(&json!({ - "model": self.model, - "state": text, - "questions": { - "matches": { "type": "noul", "instructions": statement } - } - })) - .send() - .await - .and_then(|response| response.error_for_status()) - .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))? - .json() - .await - .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?; - - let answer = &response["answers"]["matches"]; - let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p)); - match (answer["type"].as_str(), probability) { - (Some("noul"), Some(probability)) => Ok(probability), - _ => exec_err!( - "jev response must contain answers.matches with type noul and a probability in [0, 1]" - ), - } - } - - fn client(&self) -> Result { - if !self.enabled { - return exec_err!( - "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it" - ); - } - let key = self - .api_key - .as_deref() - .filter(|key| !key.trim().is_empty()) - .ok_or_else(|| { - exec_datafusion_err!("jev requires the JEV_API_KEY environment variable") - })?; - let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) - .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; - authorization.set_sensitive(true); - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert(reqwest::header::AUTHORIZATION, authorization); - Client::builder() - .default_headers(headers) - .timeout(Duration::from_secs(30)) - .build() - .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}")) - } -} - -impl Default for JevFunction { - fn default() -> Self { - Self { - // External model evaluations must not be constant-folded during planning. - signature: Signature::exact( - vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64], - Volatility::Volatile, - ), - enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), - api_key: std::env::var("JEV_API_KEY").ok(), - endpoint: std::env::var("JEV_ENDPOINT") - .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), - model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), - } - } -} - -impl fmt::Debug for JevFunction { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Credentials must never appear in plans or diagnostic output. - f.debug_struct("JevFunction") - .field("signature", &self.signature) - .field("enabled", &self.enabled) - .field("model", &self.model) - .finish_non_exhaustive() - } -} - -impl ScalarUDFImpl for JevFunction { - fn name(&self) -> &str { - "jev" - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Boolean) - } - - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { - not_impl_err!("jev can only be called from async contexts") - } -} - -#[async_trait] -impl AsyncScalarUDFImpl for JevFunction { - async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arrays = args - .args - .into_iter() - .map(|arg| arg.into_array(args.number_rows)) - .collect::>>()?; - let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?; - let texts = as_string_view_array(&texts)?; - let statements = as_string_view_array(&statements)?; - let thresholds = as_float64_array(&thresholds)?; - let rows: Vec<_> = texts - .iter() - .zip(statements.iter()) - .zip(thresholds.iter()) - .map(|((text, statement), threshold)| Some((text?, statement?, threshold?))) - .collect(); - - // Validate the whole batch before making any billable requests. - for (_, _, threshold) in rows.iter().flatten() { - if !(0.0..=1.0).contains(threshold) { - return exec_err!("jev threshold must be finite and in [0, 1]"); - } - } - if rows.iter().all(Option::is_none) { - return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( - args.number_rows, - )))); - } - - // Reuse connections within the batch; bound in-flight requests and preserve row order. - let client = self.client()?; - let requests: Vec<_> = rows - .into_iter() - .map(|row| { - let client = &client; - async move { - match row { - Some((text, statement, threshold)) => self - .evaluate(client, text, statement) - .await - .map(|p| Some(p >= threshold)), - None => Ok(None), - } - } - }) - .collect(); - // This limit is per expression/batch invocation, not per query or process. - // Concurrent partitions and queries can each have their own in-flight requests. - let matches: Vec> = stream::iter(requests).buffered(8).try_collect().await?; - Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches)))) - } -} diff --git a/src/common/function/src/scalars/jev/tests.rs b/src/common/function/src/scalars/jev/tests.rs deleted file mode 100644 index 82d7ab03013..00000000000 --- a/src/common/function/src/scalars/jev/tests.rs +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use axum::http::{HeaderMap, StatusCode}; -use axum::routing::post; -use axum::{Json, Router}; -use datafusion::assert_batches_eq; -use datafusion::prelude::{SessionConfig, SessionContext}; -use tokio::net::TcpListener; -use tokio::task::JoinHandle; - -use super::*; -use crate::function::FunctionContext; -use crate::function_registry::FUNCTION_REGISTRY; - -struct MockServer { - endpoint: String, - task: JoinHandle<()>, -} - -impl MockServer { - async fn start(router: Router) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); - Self { endpoint, task } - } - - fn function(&self) -> JevFunction { - JevFunction { - enabled: true, - api_key: Some("test-key".to_string()), - endpoint: self.endpoint.clone(), - model: "test-model".to_string(), - ..Default::default() - } - } -} - -impl Drop for MockServer { - fn drop(&mut self) { - self.task.abort(); - } -} - -fn context(function: JevFunction) -> SessionContext { - let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); - ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); - ctx -} - -#[tokio::test] -async fn test_jev_sql_filters_using_noul_probability() { - let server = MockServer::start(Router::new().route( - "/v1/systemone", - post( - |headers: HeaderMap, Json(request): Json| async move { - assert_eq!(headers["authorization"], "Bearer test-key"); - assert_eq!(request["model"], "test-model"); - assert_eq!(request["questions"]["matches"]["type"], "noul"); - // The SQL statement is passed through verbatim, not rewritten as a prompt. - assert_eq!( - request["questions"]["matches"]["instructions"], - "payment failed" - ); - let probability = match request["state"].as_str().unwrap() { - "failed" => { - // Complete this row after later rows to exercise result alignment. - tokio::time::sleep(Duration::from_millis(20)).await; - 0.95 - } - "boundary" => 0.8, - "recovered" => 0.1, - unexpected => panic!("unexpected state: {unexpected}"), - }; - Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) - }, - ), - )) - .await; - let ctx = context(server.function()); - let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)"; - let batches = ctx - .sql(&format!( - "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_batches_eq!( - ["+----+", "| id |", "+----+", "| 1 |", "| 2 |", "+----+"], - &batches - ); - - let batches = ctx - .sql(&format!( - "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_batches_eq!( - [ - "+----+---------+", - "| id | matched |", - "+----+---------+", - "| 1 | true |", - "| 2 | true |", - "| 3 | false |", - "| 4 | |", - "+----+---------+" - ], - &batches - ); -} - -#[tokio::test] -async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() { - let ctx = context(JevFunction { - enabled: true, - api_key: None, - ..Default::default() - }); - let batches = ctx - .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c") - .await.unwrap().collect().await.unwrap(); - assert_batches_eq!( - [ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| | | |", - "+---+---+---+" - ], - &batches - ); - - for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] { - let error = ctx - .sql(&format!("SELECT jev('text', 'condition', {threshold})")) - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("jev threshold must be finite and in [0, 1]"), - "{error}" - ); - } - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); - - let ctx = context(JevFunction { - enabled: false, - api_key: None, - ..Default::default() - }); - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), - "{error}" - ); -} - -#[tokio::test] -async fn test_jev_bad_api_responses_fail_the_query() { - for (status, body, expected) in [ - (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"), - (StatusCode::OK, "not json", "jev response is not valid JSON"), - ( - StatusCode::OK, - r#"{"answers":{}}"#, - "jev response must contain", - ), - ( - StatusCode::OK, - r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, - "jev response must contain", - ), - ( - StatusCode::OK, - r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, - "jev response must contain", - ), - ] { - let server = MockServer::start( - Router::new().route("/v1/systemone", post(move || async move { (status, body) })), - ) - .await; - let ctx = context(server.function()); - let error = ctx - .sql("SELECT jev('text', 'condition', 0.8)") - .await - .unwrap() - .collect() - .await - .unwrap_err(); - assert!(error.to_string().contains(expected), "{error}"); - } -} - -#[tokio::test] -#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] -async fn test_jev_live() { - let ctx = SessionContext::new(); - let function = FUNCTION_REGISTRY.get_function("jev").unwrap(); - ctx.register_udf(function.provide(FunctionContext::default())); - let batches = ctx.sql( - "SELECT id FROM (VALUES - (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), - (2, 'Payment succeeded on the second retry; the payment is complete.'), - (3, 'User logged in successfully.') - ) AS events(id, message) - WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8) - ORDER BY id", - ).await.unwrap().collect().await.unwrap(); - assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); -} diff --git a/tests/cases/standalone/common/function/ai_functions.result b/tests/cases/standalone/common/function/ai_functions.result new file mode 100644 index 00000000000..057c9ecd990 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.result @@ -0,0 +1,144 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT ai_match('message', NULL) AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + ++------------+ +| score_type | ++------------+ +| Float64 | ++------------+ + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + ++-------------+------------+ +| choice_type | score_type | ++-------------+------------+ +| Utf8 | BinaryView | ++-------------+------------+ + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + ++--------+-------+------------+------------------+ +| rating | score | confidence | high_probability | ++--------+-------+------------+------------------+ +| | | | | ++--------+-------+------------+------------------+ + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +Affected Rows: 0 + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +Affected Rows: 2 + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + ++-------------+---------+---------+ +| occurred_at | service | message | ++-------------+---------+---------+ ++-------------+---------+---------+ + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + ++---------------------+------+--------+ +| occurred_at | team | rating | ++---------------------+------+--------+ +| 2026-09-19T01:00:00 | | | +| 2026-09-19T02:00:00 | | | ++---------------------+------+--------+ + +DROP TABLE ai_events; + +Affected Rows: 0 + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_match' function: coercion from Utf8, Utf8, Float64 to the signature Exact(Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_match(Utf8, Utf8, Float64)'. You might need to add explicit type casts. + Candidate functions: + ai_match(Utf8View, Utf8View) + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria is not valid JSON: expected ident at line 1 column 2 + +SELECT ai_choose('message', 'prompt', '{}'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_choose' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_choose(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_choose(Utf8View, Utf8View, Utf8View) + +SELECT ai_score('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_score' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_score(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_score(Utf8View, Utf8View, Utf8View) + diff --git a/tests/cases/standalone/common/function/ai_functions.sql b/tests/cases/standalone/common/function/ai_functions.sql new file mode 100644 index 00000000000..fd3e80d1e18 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.sql @@ -0,0 +1,65 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + +SELECT ai_match('message', NULL) AS score; + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + +DROP TABLE ai_events; + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +SELECT ai_choose('message', 'prompt', '{}'); + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +SELECT ai_score('message', 'prompt'); diff --git a/tests/cases/standalone/common/function/jev.result b/tests/cases/standalone/common/function/jev.result deleted file mode 100644 index 4f6d91d2867..00000000000 --- a/tests/cases/standalone/common/function/jev.result +++ /dev/null @@ -1,66 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', NULL, 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -Affected Rows: 0 - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -Affected Rows: 2 - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -+-------------+---------+---------+ -| occurred_at | service | message | -+-------------+---------+---------+ -+-------------+---------+---------+ - -DROP TABLE jev_events; - -Affected Rows: 0 - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - -SELECT jev('message', 'condition', 1.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - diff --git a/tests/cases/standalone/common/function/jev.sql b/tests/cases/standalone/common/function/jev.sql deleted file mode 100644 index f4c252c9273..00000000000 --- a/tests/cases/standalone/common/function/jev.sql +++ /dev/null @@ -1,32 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -SELECT jev('message', NULL, 0.8) AS matched; - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -DROP TABLE jev_events; - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -SELECT jev('message', 'condition', 1.1); From 723da69b2182b747a82574779db32e33ffb63111 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 22 Sep 2026 14:38:17 +0000 Subject: [PATCH 59/63] feat: share logical table batching with OTLP metrics (#9288) * feat: share logical table batching with OTLP metrics Signed-off-by: WenyXu * fix: unify pending rows batch acknowledgement policy Signed-off-by: WenyXu * fix: align logical batcher example configuration expectations Signed-off-by: WenyXu * fix: align batcher worker channel defaults to 65536 Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- config/config.md | 38 ++- config/frontend.example.toml | 52 +-- config/standalone.example.toml | 52 +-- src/cmd/src/bin/query_perf_fixture/case.rs | 2 +- src/cmd/tests/load_config_test.rs | 9 + src/frontend/AGENTS.md | 5 + src/frontend/src/frontend.rs | 49 ++- src/frontend/src/instance.rs | 5 +- src/frontend/src/instance/builder.rs | 48 +-- src/frontend/src/instance/logical_batcher.rs | 102 ++++++ src/frontend/src/instance/otlp.rs | 57 +++- src/frontend/src/server.rs | 171 +++++++--- src/frontend/src/service_config.rs | 2 +- .../service_config/pending_rows_batcher.rs | 212 ++++++++++-- src/frontend/src/service_config/prom_store.rs | 2 +- src/operator/src/batcher.rs | 3 +- src/operator/src/insert.rs | 315 +++++++++++++++++- src/servers/src/batcher.rs | 18 + src/servers/src/batcher/logical_table.rs | 45 +-- .../src/batcher/logical_table/tables.rs | 44 +-- src/servers/src/batcher/table.rs | 117 ++++++- src/servers/src/batcher/table/batch.rs | 2 + src/servers/src/http.rs | 63 +++- src/session/src/context.rs | 24 +- src/standalone/src/options.rs | 73 +++- tests-integration/src/otlp.rs | 227 +++++++++++++ tests-integration/src/standalone.rs | 14 + tests-integration/tests/http.rs | 4 +- 28 files changed, 1518 insertions(+), 237 deletions(-) create mode 100644 src/frontend/src/instance/logical_batcher.rs diff --git a/config/config.md b/config/config.md index 16384b406f4..d736c24d848 100644 --- a/config/config.md +++ b/config/config.md @@ -34,7 +34,7 @@ | `runtime.experimental_workload_scheduler.sample_every_polls` | Integer | `16` | Number of polls between scheduler fairness samples. Must be greater than zero. | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -75,13 +75,20 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | -| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in HTTP ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | -| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | +| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. | +| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | @@ -94,12 +101,6 @@ | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | | `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | -| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint | -| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. | -| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. | -| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. | -| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. | -| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. | | `wal` | -- | -- | The WAL options. | | `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** | | `wal.dir` | String | Unset | The directory to store the WAL files.
**It's only used when the provider is `raft_engine`**. | @@ -289,7 +290,7 @@ | `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -342,13 +343,20 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | -| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher` | -- | -- | Ordinary-table batching for opted-in HTTP ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | -| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | +| `pending_rows_batcher.logical_table.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per physical-table worker. | +| `pending_rows_batcher.logical_table.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.logical_table.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued logical-table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | @@ -361,12 +369,6 @@ | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | | `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | -| `prom_store.pending_rows_flush_interval` | String | `0s` | Interval to flush pending rows batcher.
Set to "0s" to disable batching mode in Prometheus Remote Write endpoint | -| `prom_store.max_batch_rows` | Integer | `100000` | Max rows per pending batch before triggering a flush. | -| `prom_store.max_concurrent_flushes` | Integer | `256` | Max number of concurrent batch flushes. | -| `prom_store.worker_channel_capacity` | Integer | `65526` | Capacity of the pending batch worker channel. | -| `prom_store.max_inflight_requests` | Integer | `3000` | Max inflight write requests before backpressure. | -| `prom_store.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of logical-table flow notifications waiting in the shared queue. | | `meta_client` | -- | -- | The metasrv client options. | | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. | | `meta_client.timeout` | String | `3s` | Operation timeout. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 043da57db74..275db3105c5 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -56,9 +56,9 @@ default_column_prefix = "greptime" ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is ## raised to at least the largest active flush interval plus 1 second. The intervals come from -## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. +## `pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -231,12 +231,11 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Shared experimental ordinary-table batching for opted-in ingestion protocols. -## Legacy Prometheus batching settings under prom_store remain supported. -## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. -## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. -## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. -## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +## Ordinary-table batching for opted-in HTTP ingestion protocols. +## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge +## queue admission without waiting for storage; later failures cannot be returned to the client. +## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. +## OTLP logs, traces and ordinary metrics use this batcher. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -256,12 +255,33 @@ max_batch_rows = 100000 ## Maximum concurrent flushes shared by the frontend batcher. max_concurrent_flushes = 256 ## Maximum queued submissions per table worker. -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 ## Maximum admitted original requests awaiting completion. max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 +## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. +## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Enable independently with protocols and a nonzero flush interval. +## Omitted fields use independent defaults, not parent settings. +## Empty protocols or a zero interval disables logical batching without fallback. +## Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. +[pending_rows_batcher.logical_table] +# protocols = ["prom", "otlp"] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by Prom and OTLP metrics. +max_concurrent_flushes = 256 +## Maximum queued submissions per physical-table worker. +worker_channel_capacity = 65536 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued logical-table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. @@ -293,19 +313,7 @@ with_metric_engine = true prom_validation_mode = "strict" ## Experimental: enable Prometheus remote write v2 native histogram ingestion. experimental_enable_prometheus_native_histogram = false -## Interval to flush pending rows batcher. -## Set to "0s" to disable batching mode in Prometheus Remote Write endpoint -#+pending_rows_flush_interval = "0s" -## Max rows per pending batch before triggering a flush. -#+max_batch_rows = 100000 -## Max number of concurrent batch flushes. -#+max_concurrent_flushes = 256 -## Capacity of the pending batch worker channel. -#+worker_channel_capacity = 65526 -## Max inflight write requests before backpressure. -#+max_inflight_requests = 3000 -## Maximum number of logical-table flow notifications waiting in the shared queue. -#+flow_notification_queue_capacity = 1024 + ## The metasrv client options. [meta_client] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e59e7ba1b4c..257f48eda2c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -81,9 +81,9 @@ max_concurrent_queries = 0 ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## When synchronous Prometheus, OTLP metrics, or ordinary-table batching is enabled, a nonzero timeout is ## raised to at least the largest active flush interval plus 1 second. The intervals come from -## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. +## `pending_rows_batcher.logical_table.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -210,12 +210,11 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Shared experimental ordinary-table batching for opted-in ingestion protocols. -## Legacy Prometheus batching settings under prom_store remain supported. -## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. -## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. -## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. -## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +## Ordinary-table batching for opted-in HTTP ingestion protocols. +## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge +## queue admission without waiting for storage; later failures cannot be returned to the client. +## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. +## OTLP logs, traces and ordinary metrics use this batcher. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -235,12 +234,33 @@ max_batch_rows = 100000 ## Maximum concurrent flushes shared by the frontend batcher. max_concurrent_flushes = 256 ## Maximum queued submissions per table worker. -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 ## Maximum admitted original requests awaiting completion. max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 +## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. +## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Enable independently with protocols and a nonzero flush interval. +## Omitted fields use independent defaults, not parent settings. +## Empty protocols or a zero interval disables logical batching without fallback. +## Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. +[pending_rows_batcher.logical_table] +# protocols = ["prom", "otlp"] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by Prom and OTLP metrics. +max_concurrent_flushes = 256 +## Maximum queued submissions per physical-table worker. +worker_channel_capacity = 65536 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued logical-table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. @@ -272,19 +292,7 @@ with_metric_engine = true prom_validation_mode = "strict" ## Experimental: enable Prometheus remote write v2 native histogram ingestion. experimental_enable_prometheus_native_histogram = false -## Interval to flush pending rows batcher. -## Set to "0s" to disable batching mode in Prometheus Remote Write endpoint -#+pending_rows_flush_interval = "0s" -## Max rows per pending batch before triggering a flush. -#+max_batch_rows = 100000 -## Max number of concurrent batch flushes. -#+max_concurrent_flushes = 256 -## Capacity of the pending batch worker channel. -#+worker_channel_capacity = 65526 -## Max inflight write requests before backpressure. -#+max_inflight_requests = 3000 -## Maximum number of logical-table flow notifications waiting in the shared queue. -#+flow_notification_queue_capacity = 1024 + ## The WAL options. [wal] diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs index 0f5239a6b51..271c3928e35 100644 --- a/src/cmd/src/bin/query_perf_fixture/case.rs +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -185,7 +185,7 @@ pub(super) fn default_max_concurrent_flushes() -> u64 { 256 } pub(super) fn default_worker_channel_capacity() -> u64 { - 65526 + 65_536 } pub(super) fn default_max_inflight_requests() -> u64 { 3000 diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 67e2e573f38..1c4a83e1397 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -30,6 +30,7 @@ use datanode::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use file_engine::config::EngineConfig as FileEngineConfig; use flow::FlownodeOptions; use frontend::frontend::FrontendOptions; +use frontend::service_config::PendingRowsBatcherOptions; use meta_client::MetaClientOptions; use meta_srv::metasrv::MetasrvOptions; use meta_srv::selector::SelectorType; @@ -206,6 +207,10 @@ fn test_load_frontend_example_config() { ); let expected = GreptimeOptions:: { component: FrontendOptions { + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: Some(Default::default()), + ..Default::default() + }, default_timezone: Some("UTC".to_string()), default_column_prefix: Some("greptime".to_string()), auto_create_table: true, @@ -389,6 +394,10 @@ fn test_load_standalone_example_config() { ); let expected = GreptimeOptions:: { component: StandaloneOptions { + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: Some(Default::default()), + ..Default::default() + }, default_timezone: Some("UTC".to_string()), default_column_prefix: Some("greptime".to_string()), auto_create_table: true, diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index d552d73cf39..49a4b8fcd0f 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -51,6 +51,11 @@ remote datanodes via `operator`/`client`. - Internal gRPC listeners mark requests with `Channel::Internal` in middleware (`server.rs`), including requests handled by Enterprise Flight wrappers. +- **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes + one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine + writes. OTLP checks operator eligibility and falls back for incompatible tables. + The schema adapter holds a weak instance reference to avoid an ownership cycle. + - **Table batching** (`instance/builder.rs`): protocol entry points opt in through `QueryContext`. The primary inserter prepares eligible ordinary-table writes for `servers::batcher::table::TablePendingRowsBatcher`. A separate execution-only diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 86ba2cb9bfb..72bf322b1f1 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -67,7 +67,7 @@ pub struct FrontendOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, - /// Shared experimental ordinary-table batching; independent of Prom batching. + /// Ordinary-table batching with independent logical-table controls. pub pending_rows_batcher: PendingRowsBatcherOptions, pub prom_store: PromStoreOptions, pub jaeger: JaegerOptions, @@ -131,6 +131,7 @@ impl Configurable for FrontendOptions { "meta_client.metasrv_addrs", "event_recorder.event_types", "pending_rows_batcher.protocols", + "pending_rows_batcher.logical_table.protocols", ]) } } @@ -225,6 +226,38 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_logical_batcher_protocols_from_env() { + temp_env::with_vars( + [ + ( + "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("otlp,influxdb"), + ), + ( + "FRONTEND_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS", + Some("prom,otlp"), + ), + ], + || { + let options = + FrontendOptions::load_layered_options(None, "FRONTEND_LOGICAL_TEST").unwrap(); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .protocols, + vec![ + servers::http::BatchingProtocol::Prom, + servers::http::BatchingProtocol::Otlp + ] + ); + }, + ); + } + #[test] fn test_batcher_protocols_from_env() { temp_env::with_vars( @@ -236,7 +269,7 @@ mod tests { let options = FrontendOptions::load_layered_options(None, "FRONTEND_BATCHER_TEST").unwrap(); assert_eq!( - options.pending_rows_batcher.protocols, + options.pending_rows_batcher.table.protocols, vec![ servers::http::BatchingProtocol::Influxdb, servers::http::BatchingProtocol::HttpSql @@ -252,6 +285,7 @@ mod tests { assert!( !defaults .pending_rows_batcher + .table .pending_rows_batching_enabled() ); let options: FrontendOptions = toml::from_str( @@ -263,9 +297,14 @@ max_batch_rows = 25 "#, ) .unwrap(); - assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); - assert_eq!(options.pending_rows_batcher.protocols.len(), 2); - assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + assert_eq!(options.pending_rows_batcher.table.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert!( + options + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); let serialized = toml::to_string(&options).unwrap(); let parsed: FrontendOptions = toml::from_str(&serialized).unwrap(); assert_eq!(options.influxdb, parsed.influxdb); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 22d2afaa7bc..92bc74ca10b 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -21,6 +21,7 @@ mod import_packed; mod influxdb; mod jaeger; mod log_handler; +mod logical_batcher; mod logs; mod opentsdb; mod otlp; @@ -31,7 +32,7 @@ mod region_query; use std::collections::HashSet; use std::pin::Pin; use std::sync::atomic::AtomicBool; -use std::sync::{Arc, atomic}; +use std::sync::{Arc, OnceLock, atomic}; use std::time::{Duration, SystemTime}; use async_stream::stream; @@ -78,6 +79,7 @@ use query::metrics::OnDone; use query::parser::{PromQuery, QueryStatement}; use query::query_engine::DescribeResult; use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; +use servers::batcher::logical_table::LogicalTablePendingRowsBatcher; use servers::error::{ self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu, OtlpMetricModeIncompatibleSnafu, UnexpectedResultSnafu, @@ -130,6 +132,7 @@ pub struct Instance { query_engine: QueryEngineRef, plugins: Plugins, inserter: InserterRef, + logical_batcher: Arc>>>, deleter: DeleterRef, table_metadata_manager: TableMetadataManagerRef, event_recorder: EventRecorderRef, diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 01e9ef034e3..a3cfbcdaffe 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -61,7 +61,7 @@ use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; use crate::instance::entity_graph::EntityGraphProviderImpl; use crate::instance::region_query::FrontendRegionQueryHandler; -use crate::service_config::PendingRowsBatcherOptions; +use crate::service_config::BatcherOptions; /// The frontend [`Instance`] builder. pub struct FrontendBuilder { @@ -237,31 +237,30 @@ impl FrontendBuilder { }; // The execution-only inserter owns no batchers, avoiding an Arc cycle. let bulk_inserter = Arc::new(create_inserter()); - let build_batcher = - |options: &PendingRowsBatcherOptions| -> Option> { - if !options.pending_rows_batching_enabled() - || (self.options.prom_store.with_metric_engine - && options - .protocols - .iter() - .all(|protocol| *protocol == BatchingProtocol::Prom)) - { - return None; - } - TablePendingRowsBatcher::try_new( - options.pending_rows_flush_interval, - options.max_batch_rows, - options.max_concurrent_flushes, - options.worker_channel_capacity, - options.max_inflight_requests, - options.flow_notification_queue_capacity, - bulk_inserter.clone(), - ) - .map(|batcher| batcher as Arc) - }; + let build_batcher = |options: &BatcherOptions| -> Option> { + if !options.pending_rows_batching_enabled() + || (self.options.prom_store.with_metric_engine + && options + .protocols + .iter() + .all(|protocol| *protocol == BatchingProtocol::Prom)) + { + return None; + } + TablePendingRowsBatcher::try_new( + options.pending_rows_flush_interval, + options.max_batch_rows, + options.max_concurrent_flushes, + options.worker_channel_capacity, + options.max_inflight_requests, + options.flow_notification_queue_capacity, + bulk_inserter.clone(), + ) + .map(|batcher| batcher as Arc) + }; let inserter = Arc::new( create_inserter() - .with_pending_rows_batcher(build_batcher(&self.options.pending_rows_batcher)), + .with_pending_rows_batcher(build_batcher(self.options.table_batcher_options())), ); let deleter = Arc::new(Deleter::new( self.catalog_manager.clone(), @@ -386,6 +385,7 @@ impl FrontendBuilder { admin_event_recorder.install(&event_recorder); Ok(Instance { + logical_batcher: Default::default(), frontend_peer_addr, experimental_metric_export: self.options.experimental_metric_export, catalog_manager: self.catalog_manager, diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs new file mode 100644 index 00000000000..28ed19b3e14 --- /dev/null +++ b/src/frontend/src/instance/logical_batcher.rs @@ -0,0 +1,102 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, Weak}; + +use api::v1::ColumnSchema; +use async_trait::async_trait; +use servers::batcher::logical_table::{LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer}; +use servers::error::{BatcherChannelClosedSnafu, Result}; +use servers::http::BatchingProtocol; +use session::context::QueryContextRef; +use snafu::OptionExt; + +use crate::frontend::FrontendOptions; +use crate::instance::Instance; + +impl Instance { + pub(crate) fn init_logical_batcher(self: &Arc, options: &FrontendOptions) { + self.logical_batcher.get_or_init(|| { + let options_batcher = options.logical_batcher_options(); + let enabled = options_batcher + .protocols + .iter() + .any(|protocol| match protocol { + BatchingProtocol::Prom => options.prom_store.enable, + BatchingProtocol::Otlp => options.otlp.enable, + _ => false, + }); + if !options.prom_store.with_metric_engine + || !enabled + || !options_batcher.pending_rows_batching_enabled() + { + return None; + } + LogicalTablePendingRowsBatcher::try_new( + self.partition_manager().clone(), + self.node_manager().clone(), + self.catalog_manager().clone(), + self.table_flownode_set_cache().clone(), + true, + Arc::new(LogicalTables(Arc::downgrade(self))), + options_batcher.pending_rows_flush_interval, + options_batcher.max_batch_rows, + options_batcher.max_concurrent_flushes, + options_batcher.worker_channel_capacity, + options_batcher.max_inflight_requests, + options_batcher.flow_notification_queue_capacity, + ) + }); + } + + pub(crate) fn logical_batcher(&self) -> Option<&Arc> { + self.logical_batcher.get().and_then(Option::as_ref) + } +} + +// The instance owns the shared batcher; schema preparation must not keep that +// owner alive through a reference cycle. +struct LogicalTables(Weak); + +#[async_trait] +impl PendingRowsSchemaAlterer for LogicalTables { + async fn create_tables_if_missing_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[ColumnSchema])], + with_metric_engine: bool, + ctx: QueryContextRef, + ) -> Result<()> { + self.0 + .upgrade() + .context(BatcherChannelClosedSnafu)? + .create_tables_if_missing_batch(catalog, schema, tables, with_metric_engine, ctx) + .await + } + + async fn add_missing_prom_tag_columns_batch( + &self, + catalog: &str, + schema: &str, + tables: &[(&str, &[String])], + ctx: QueryContextRef, + ) -> Result<()> { + self.0 + .upgrade() + .context(BatcherChannelClosedSnafu)? + .add_missing_prom_tag_columns_batch(catalog, schema, tables, ctx) + .await + } +} diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index fa40c7eb59a..7a0f706c58d 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -27,10 +27,11 @@ use client::Output; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::{OutputData, OutputMeta}; use common_telemetry::{tracing, warn}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; -use operator::insert::{admit_row_insert_batches, admit_write}; +use operator::insert::{Inserter, admit_row_insert_batches, admit_write}; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, AuthSnafu, Result as ServerResult}; @@ -165,27 +166,59 @@ impl OpenTelemetryProtocolHandler for Instance { Arc::new(c) }; + let physical_table = ctx + .extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .to_string(); + let batcher = self.logical_batcher().filter(|_| { + ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine + }); + let batcher = if batcher.is_some() + && self + .inserter + .can_batch_metric_rows(&requests, &ctx, &physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)? + { + batcher + } else { + None + }; + // OTLP tables have one sample field in both the legacy and physical paths. - let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + let output = if let Some(batcher) = batcher { + let (rows, cost) = batcher + .submit_with(requests, ctx.clone(), |mut requests| { + let ctx = ctx.clone(); + async move { + Inserter::meter_row_inserts(&mut requests, &ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + }) + .await?; + Output::new( + OutputData::AffectedRows(rows as usize), + OutputMeta::new_with_cost(cost as _), + ) + } else if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) + .context(error::ExecuteGrpcQuerySnafu)? } else { - let physical_table = ctx - .extension(PHYSICAL_TABLE_PARAM) - .unwrap_or(GREPTIME_PHYSICAL_TABLE) - .to_string(); self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) .await .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - }?; + .context(error::ExecuteGrpcQuerySnafu)? + }; outcome.write_cost = output.meta.cost; - // Derived enrichment, written after the metric data is committed: - // failing here would make the client retry data the server already - // accepted, so every failure degrades to a warning instead. + // Derived enrichment follows the accepted metric submission, which may + // still be queued in asynchronous mode. Failures remain warning-only + // to avoid retrying metric data the server already accepted. if let Some(resource_info) = resource_info { let written = match self.check_row_insert_permission( &resource_info, diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 017bba53fba..0a6612ba31c 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -24,9 +24,7 @@ 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::batcher::pending_rows_batch_sync_enabled; use servers::error::Error as ServerError; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::FlightCraftRef; @@ -74,6 +72,7 @@ where { pub fn new(opts: T, instance: Arc, plugins: Plugins) -> Self { let feopts = opts.clone().into(); + instance.init_logical_batcher(&feopts); // Create server request memory limiter for all server protocols let server_memory_limiter = ServerMemoryLimiter::new( feopts.max_in_flight_write_bytes.as_bytes(), @@ -110,7 +109,8 @@ where request_memory_limiter: ServerMemoryLimiter, ) -> HttpServerBuilder { let mut builder = HttpServerBuilder::new(effective_http_options(opts)) - .with_batching_protocols(opts.pending_rows_batcher.protocols.clone()) + .with_batching_protocols(opts.table_batcher_options().protocols.clone()) + .with_logical_batching_protocols(opts.logical_batcher_options().protocols) .with_memory_limiter(request_memory_limiter) .with_sql_handler(self.instance.clone()); @@ -134,24 +134,12 @@ where let prom_store = effective_prom_store_options(opts); if prom_store.enable { - let pending_rows_batcher = if prom_store.with_metric_engine { - LogicalTablePendingRowsBatcher::try_new( - self.instance.partition_manager().clone(), - self.instance.node_manager().clone(), - self.instance.catalog_manager().clone(), - self.instance.table_flownode_set_cache().clone(), - prom_store.with_metric_engine, - self.instance.clone(), - prom_store.pending_rows_flush_interval, - prom_store.max_batch_rows, - prom_store.max_concurrent_flushes, - prom_store.worker_channel_capacity, - prom_store.max_inflight_requests, - prom_store.flow_notification_queue_capacity, - ) - } else { - None - }; + let pending_rows_batcher = opts + .logical_batcher_options() + .protocols + .contains(&BatchingProtocol::Prom) + .then(|| self.instance.logical_batcher().cloned()) + .flatten(); builder = builder .with_prom_handler( self.instance.clone(), @@ -439,15 +427,16 @@ where /// Selected shared controls override legacy Prom batching knobs, not protocol behavior. fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions { let mut prom_store = opts.prom_store.clone(); - let shared = &opts.pending_rows_batcher; - if shared.protocols.contains(&BatchingProtocol::Prom) && shared.pending_rows_batching_enabled() - { + let shared = opts.logical_batcher_options(); + if shared.protocols.contains(&BatchingProtocol::Prom) { prom_store.pending_rows_flush_interval = shared.pending_rows_flush_interval; prom_store.max_batch_rows = shared.max_batch_rows; prom_store.max_concurrent_flushes = shared.max_concurrent_flushes; prom_store.worker_channel_capacity = shared.worker_channel_capacity; prom_store.max_inflight_requests = shared.max_inflight_requests; prom_store.flow_notification_queue_capacity = shared.flow_notification_queue_capacity; + } else { + prom_store.pending_rows_flush_interval = Duration::ZERO; } prom_store } @@ -459,10 +448,9 @@ fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions { let mut http = opts.http.clone(); let prom_store = effective_prom_store_options(opts); - let shared = &opts.pending_rows_batcher; - // Ordinary-table batching always waits for its flush, independently of the - // dedicated Prom batcher's asynchronous acknowledgement mode. - let common_enabled = shared.pending_rows_batching_enabled() + let shared = opts.table_batcher_options(); + let common_enabled = batch_sync + && shared.pending_rows_batching_enabled() && shared.protocols.iter().any(|protocol| { *protocol != BatchingProtocol::Prom || (prom_store.enable && !prom_store.with_metric_engine) @@ -470,7 +458,19 @@ fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval); let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync) .then_some(prom_store.pending_rows_flush_interval); - let Some(flush_interval) = common_interval.into_iter().chain(prom_interval).max() else { + let logical = opts.logical_batcher_options(); + let otlp_interval = (batch_sync + && opts.otlp.enable + && opts.prom_store.with_metric_engine + && logical.protocols.contains(&BatchingProtocol::Otlp) + && logical.pending_rows_batching_enabled()) + .then_some(logical.pending_rows_flush_interval); + let Some(flush_interval) = common_interval + .into_iter() + .chain(prom_interval) + .chain(otlp_interval) + .max() + else { return http; }; let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1)); @@ -515,6 +515,81 @@ mod tests { use crate::instance::builder::FrontendBuilder; use crate::server::*; + #[tokio::test] + async fn test_logical_batcher_shared_without_prom_endpoint() { + use crate::service_config::BatcherOptions; + let mut options = FrontendOptions::default(); + options.prom_store.enable = false; + options.pending_rows_batcher.logical_table = Some(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }); + let meta_client = Arc::new( + MetaClientBuilder::new(0, Role::Frontend) + .enable_procedure() + .build(), + ); + let instance = Arc::new( + FrontendBuilder::new_test(&options, meta_client) + .try_build() + .await + .unwrap(), + ); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let batcher = instance.logical_batcher().unwrap().clone(); + instance.init_logical_batcher(&options); + assert!(Arc::ptr_eq(&batcher, instance.logical_batcher().unwrap())); + let weak = Arc::downgrade(&instance); + drop(services); + drop(instance); + assert!( + weak.upgrade().is_none(), + "batcher must not retain its schema owner" + ); + } + + #[test] + fn test_logical_batcher_http_timeout_and_prom_disable() { + use crate::service_config::pending_rows_batcher::BatcherOptions; + let mut opts = FrontendOptions::default(); + opts.http.timeout = Duration::from_secs(1); + opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2); + opts.pending_rows_batcher.logical_table = Some(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(5), + ..Default::default() + }); + assert!(!effective_prom_store_options(&opts).pending_rows_batching_enabled()); + // Both logical protocols follow the global acknowledgement policy, + // independently of enabling the Prom HTTP endpoint. + opts.prom_store.enable = false; + assert_eq!( + effective_http_options_with_sync(&opts, true).timeout, + Duration::from_secs(6) + ); + assert_eq!( + effective_http_options_with_sync(&opts, false).timeout, + opts.http.timeout + ); + opts.prom_store.with_metric_engine = false; + assert_eq!( + effective_http_options_with_sync(&opts, false).timeout, + opts.http.timeout + ); + opts.prom_store.with_metric_engine = true; + opts.pending_rows_batcher + .logical_table + .as_mut() + .unwrap() + .protocols + .clear(); + assert_eq!( + effective_http_options_with_sync(&opts, true).timeout, + opts.http.timeout + ); + } + #[test] fn test_effective_prom_batching_controls() { // Only an enabled shared Prom selection replaces the legacy controls. @@ -532,7 +607,7 @@ mod tests { opts.prom_store.enable = prom_enabled; opts.prom_store .experimental_enable_prometheus_native_histogram = true; - let shared = &mut opts.pending_rows_batcher; + let shared = &mut opts.pending_rows_batcher.table; shared.protocols = vec![if selected { BatchingProtocol::Prom } else { @@ -569,7 +644,7 @@ mod tests { #[test] fn test_http_timeout_covers_synchronous_batchers() { - // Shared ordinary writes remain synchronous even when Prom is asynchronous. + // Only synchronous batchers extend the HTTP timeout. for ( protocols, metric_engine, @@ -579,10 +654,10 @@ mod tests { timeout_secs, expected_secs, ) in [ - (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 1), (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1), (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6), - (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 1), (vec![BatchingProtocol::Influxdb], true, true, 5, 8, 1, 9), (vec![BatchingProtocol::Influxdb], true, true, 8, 5, 1, 9), (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 0, 0), @@ -595,8 +670,8 @@ mod tests { opts.http.timeout = Duration::from_secs(timeout_secs); opts.prom_store.with_metric_engine = metric_engine; opts.prom_store.pending_rows_flush_interval = Duration::from_secs(legacy_secs); - opts.pending_rows_batcher.protocols = protocols; - opts.pending_rows_batcher.pending_rows_flush_interval = + opts.pending_rows_batcher.table.protocols = protocols; + opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(shared_secs); assert_eq!( effective_http_options_with_sync(&opts, batch_sync).timeout, @@ -697,24 +772,30 @@ mod tests { fn test_invalid_shared_batching_preserves_prom_store_options() { type KnobMutator = fn(&mut FrontendOptions); let cases: [KnobMutator; 5] = [ - |opts| opts.pending_rows_batcher.max_concurrent_flushes = usize::MAX, - |opts| opts.pending_rows_batcher.worker_channel_capacity = usize::MAX, - |opts| opts.pending_rows_batcher.max_inflight_requests = usize::MAX, + |opts| opts.pending_rows_batcher.table.max_concurrent_flushes = usize::MAX, + |opts| opts.pending_rows_batcher.table.worker_channel_capacity = usize::MAX, + |opts| opts.pending_rows_batcher.table.max_inflight_requests = usize::MAX, |opts| { - opts.pending_rows_batcher.flow_notification_queue_capacity = - NonZeroUsize::new(usize::MAX).unwrap() + opts.pending_rows_batcher + .table + .flow_notification_queue_capacity = NonZeroUsize::new(usize::MAX).unwrap() }, - |opts| opts.pending_rows_batcher.pending_rows_flush_interval = Duration::MAX, + |opts| opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::MAX, ]; for invalidate in cases { let mut opts = FrontendOptions::default(); opts.http.timeout = Duration::from_secs(1); opts.prom_store.pending_rows_flush_interval = Duration::from_secs(5); - opts.pending_rows_batcher.protocols = + opts.pending_rows_batcher.table.protocols = vec![BatchingProtocol::Prom, BatchingProtocol::Influxdb]; - opts.pending_rows_batcher.pending_rows_flush_interval = Duration::from_secs(10); + opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(10); invalidate(&mut opts); - assert!(!opts.pending_rows_batcher.pending_rows_batching_enabled()); + assert!( + !opts + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); assert_eq!(opts.prom_store, effective_prom_store_options(&opts)); assert_eq!( Duration::from_secs(6), diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs index bcd1097cfcf..2c761c4242a 100644 --- a/src/frontend/src/service_config.rs +++ b/src/frontend/src/service_config.rs @@ -26,6 +26,6 @@ pub use jaeger::JaegerOptions; pub use mysql::MysqlOptions; pub use opentsdb::OpentsdbOptions; pub use otlp::OtlpOptions; -pub use pending_rows_batcher::PendingRowsBatcherOptions; +pub use pending_rows_batcher::{BatcherOptions, PendingRowsBatcherOptions}; pub use postgres::PostgresOptions; pub use prom_store::PromStoreOptions; diff --git a/src/frontend/src/service_config/pending_rows_batcher.rs b/src/frontend/src/service_config/pending_rows_batcher.rs index ef7ce14ef02..61572015d7c 100644 --- a/src/frontend/src/service_config/pending_rows_batcher.rs +++ b/src/frontend/src/service_config/pending_rows_batcher.rs @@ -20,10 +20,28 @@ use serde::{Deserialize, Serialize}; use servers::http::BatchingProtocol; use tokio::sync::Semaphore; -/// Experimental table write batching options shared by HTTP ingestion protocols. -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +use crate::frontend::FrontendOptions; + +/// Independent ordinary-table and logical-table batching configuration. +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct PendingRowsBatcherOptions { + /// Existing ordinary-table controls retain their original TOML paths. + #[serde(flatten)] + pub table: BatcherOptions, + /// Independent logical-table controls; absence retains the legacy Prom fallback. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_logical_options" + )] + pub logical_table: Option, +} + +/// Write batching controls shared by HTTP ingestion protocols. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default)] +pub struct BatcherOptions { /// HTTP write protocols sharing this batcher; empty disables all entrances. pub protocols: Vec, /// Time from the first pending submission to a timed flush. Zero disables batching. @@ -41,7 +59,7 @@ pub struct PendingRowsBatcherOptions { pub flow_notification_queue_capacity: NonZeroUsize, } -impl PendingRowsBatcherOptions { +impl BatcherOptions { /// Returns whether a protocol opts in and its controls pass construction validation. pub fn pending_rows_batching_enabled(&self) -> bool { !self.protocols.is_empty() @@ -58,79 +76,229 @@ impl PendingRowsBatcherOptions { } } -impl Default for PendingRowsBatcherOptions { +impl Default for BatcherOptions { fn default() -> Self { Self { protocols: Vec::new(), pending_rows_flush_interval: Duration::ZERO, max_batch_rows: 100_000, max_concurrent_flushes: 256, - worker_channel_capacity: 65526, + worker_channel_capacity: 65_536, max_inflight_requests: 3000, flow_notification_queue_capacity: NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN), } } } +/// Rejects protocols that cannot write metric-engine logical tables. +fn deserialize_logical_options<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let options = Option::::deserialize(deserializer)?; + if options.as_ref().is_some_and(|options| { + options + .protocols + .iter() + .any(|protocol| !matches!(protocol, BatchingProtocol::Prom | BatchingProtocol::Otlp)) + }) { + return Err(serde::de::Error::custom( + "pending_rows_batcher.logical_table only supports prom and otlp", + )); + } + Ok(options) +} + +impl FrontendOptions { + /// Explicit configuration wins even when it disables batching. + pub(crate) fn table_batcher_options(&self) -> &BatcherOptions { + &self.pending_rows_batcher.table + } + + /// Resolve legacy Prom controls once without merging fields across config blocks. + pub(crate) fn logical_batcher_options(&self) -> BatcherOptions { + if let Some(options) = &self.pending_rows_batcher.logical_table { + return options.clone(); + } + let shared = &self.pending_rows_batcher.table; + if shared.protocols.contains(&BatchingProtocol::Prom) + && shared.pending_rows_batching_enabled() + { + let mut options = shared.clone(); + options.protocols = vec![BatchingProtocol::Prom]; + return options; + } + let prom = &self.prom_store; + BatcherOptions { + protocols: vec![BatchingProtocol::Prom], + pending_rows_flush_interval: prom.pending_rows_flush_interval, + max_batch_rows: prom.max_batch_rows, + max_concurrent_flushes: prom.max_concurrent_flushes, + worker_channel_capacity: prom.worker_channel_capacity, + max_inflight_requests: prom.max_inflight_requests, + flow_notification_queue_capacity: prom.flow_notification_queue_capacity, + } + } +} + #[cfg(test)] mod tests { use crate::service_config::pending_rows_batcher::*; + #[test] + fn test_batcher_config_precedence() { + use crate::frontend::FrontendOptions; + let legacy = "[prom_store]\nenable = true\nwith_metric_engine = true\npending_rows_flush_interval = '2s'\nmax_batch_rows = 9\n"; + let shared = "[pending_rows_batcher]\nprotocols = ['prom', 'otlp']\npending_rows_flush_interval = '3s'\n"; + let options: FrontendOptions = toml::from_str(legacy).unwrap(); + assert_eq!(options.logical_batcher_options().max_batch_rows, 9); + assert_eq!( + options.logical_batcher_options().protocols, + vec![BatchingProtocol::Prom] + ); + let options: FrontendOptions = toml::from_str(&format!("{legacy}{shared}")).unwrap(); + assert_eq!( + options + .logical_batcher_options() + .pending_rows_flush_interval, + Duration::from_secs(3) + ); + assert_eq!(options.logical_batcher_options().max_batch_rows, 100_000); + assert_eq!( + options.logical_batcher_options().protocols, + vec![BatchingProtocol::Prom] + ); + for explicit in [ + "", + "protocols = ['otlp']\npending_rows_flush_interval = '5s'", + ] { + let options: FrontendOptions = toml::from_str(&format!( + "{legacy}{shared}[pending_rows_batcher.logical_table]\n{explicit}" + )) + .unwrap(); + assert_eq!( + &options.logical_batcher_options(), + options.pending_rows_batcher.logical_table.as_ref().unwrap() + ); + assert_eq!( + options.table_batcher_options(), + &toml::from_str::(&format!("{legacy}{shared}")) + .unwrap() + .pending_rows_batcher + .table + ); + let restored: FrontendOptions = + toml::from_str(&toml::to_string(&options).unwrap()).unwrap(); + assert_eq!( + options.pending_rows_batcher.logical_table, + restored.pending_rows_batcher.logical_table + ); + assert_eq!(options.pending_rows_batcher, restored.pending_rows_batcher); + } + } + + #[test] + fn test_batcher_config_independent_defaults() { + use crate::frontend::FrontendOptions; + + let parent = "[pending_rows_batcher]\nprotocols = ['prom']\npending_rows_flush_interval = '3s'\nmax_batch_rows = 7\n"; + let options: FrontendOptions = toml::from_str(parent).unwrap(); + let serialized = toml::to_string(&options).unwrap(); + assert!(!serialized.contains("[pending_rows_batcher.logical_table]")); + assert!(!serialized.contains("[pending_rows_batcher.table]")); + let restored: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(restored.logical_batcher_options().max_batch_rows, 7); + + let options: FrontendOptions = + toml::from_str(&format!("{parent}[pending_rows_batcher.logical_table]")).unwrap(); + assert_eq!(options.table_batcher_options().max_batch_rows, 7); + assert_eq!(options.logical_batcher_options(), BatcherOptions::default()); + assert!( + !options + .logical_batcher_options() + .pending_rows_batching_enabled() + ); + } + + #[test] + fn test_logical_batcher_config_protocols() { + use crate::frontend::FrontendOptions; + for protocols in ["[]", "['prom']", "['otlp']", "['prom', 'otlp']"] { + let options: FrontendOptions = toml::from_str(&format!( + "[pending_rows_batcher.logical_table]\nprotocols = {protocols}" + )) + .unwrap(); + assert!(options.pending_rows_batcher.logical_table.is_some()); + } + for protocol in [ + "influxdb", + "logs", + "loki", + "http_sql", + "opentsdb", + "elasticsearch", + "splunk", + ] { + assert!( + toml::from_str::(&format!( + "[pending_rows_batcher.logical_table]\nprotocols = ['{protocol}']" + )) + .is_err() + ); + } + } + #[test] fn test_protocols() { - let options: PendingRowsBatcherOptions = toml::from_str( + let options: BatcherOptions = toml::from_str( "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom']", ).unwrap(); assert_eq!(options.protocols.len(), 9); assert!(options.protocols.contains(&BatchingProtocol::HttpSql)); - assert!(PendingRowsBatcherOptions::default().protocols.is_empty()); + assert!(BatcherOptions::default().protocols.is_empty()); for invalid in ["sql", "jaeger", "unknown"] { assert!( - toml::from_str::(&format!("protocols = ['{invalid}']")) - .is_err() + toml::from_str::(&format!("protocols = ['{invalid}']")).is_err() ); } } #[test] fn test_notification_capacity() { - let default = PendingRowsBatcherOptions::default(); + let default = BatcherOptions::default(); assert_eq!(default.flow_notification_queue_capacity.get(), 1024); - let configured: PendingRowsBatcherOptions = + let configured: BatcherOptions = toml::from_str("flow_notification_queue_capacity = 8").unwrap(); assert_eq!(configured.flow_notification_queue_capacity.get(), 8); - assert!( - toml::from_str::("flow_notification_queue_capacity = 0") - .is_err() - ); + assert!(toml::from_str::("flow_notification_queue_capacity = 0").is_err()); } #[test] fn test_defaults_and_roundtrip() { - let options: PendingRowsBatcherOptions = toml::from_str("").unwrap(); - assert_eq!(options, PendingRowsBatcherOptions::default()); + let options: BatcherOptions = toml::from_str("").unwrap(); + assert_eq!(options, BatcherOptions::default()); assert!(!options.pending_rows_batching_enabled()); assert_eq!(options.max_batch_rows, 100_000); assert_eq!(options.max_concurrent_flushes, 256); - assert_eq!(options.worker_channel_capacity, 65526); + assert_eq!(options.worker_channel_capacity, 65_536); assert_eq!(options.max_inflight_requests, 3000); let serialized = toml::to_string(&options).unwrap(); assert_eq!( options, - toml::from_str::(&serialized).unwrap() + toml::from_str::(&serialized).unwrap() ); } #[test] fn test_partial_options_and_zero_controls() { - let options: PendingRowsBatcherOptions = + let options: BatcherOptions = toml::from_str("pending_rows_flush_interval = '5ms'").unwrap(); assert_eq!( options.pending_rows_flush_interval, Duration::from_millis(5) ); assert!(!options.pending_rows_batching_enabled()); - let enabled: PendingRowsBatcherOptions = + let enabled: BatcherOptions = toml::from_str("protocols = ['http_sql']\npending_rows_flush_interval = '5ms'") .unwrap(); assert!(enabled.pending_rows_batching_enabled()); @@ -140,7 +308,7 @@ mod tests { "worker_channel_capacity", "max_inflight_requests", ] { - let options: PendingRowsBatcherOptions = toml::from_str(&format!( + let options: BatcherOptions = toml::from_str(&format!( "protocols = ['influxdb']\npending_rows_flush_interval = '5ms'\n{field} = 0" )) .unwrap(); diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs index f8429a86925..2fdd05043db 100644 --- a/src/frontend/src/service_config/prom_store.rs +++ b/src/frontend/src/service_config/prom_store.rs @@ -54,7 +54,7 @@ fn default_max_concurrent_flushes() -> usize { } fn default_worker_channel_capacity() -> usize { - 65526 + 65_536 } fn default_max_inflight_requests() -> usize { diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs index 4c5e9892c0b..677524d53f3 100644 --- a/src/operator/src/batcher.rs +++ b/src/operator/src/batcher.rs @@ -29,7 +29,8 @@ pub trait PendingRowsBatcher: Send + Sync { /// Acquires one slot per original request, shared by all of its table submissions. async fn acquire(&self) -> Result>; - /// Waits for the submitted rows to be written, retaining the slot through completion. + /// Submits rows according to the acknowledgement policy, retaining the slot until + /// writing completes even when the response acknowledges queue admission only. /// Cancelling the response wait does not retract an already enqueued submission. async fn submit( &self, diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index afc903621e2..49215f89713 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -156,6 +156,127 @@ pub struct InstantAndNormalInsertRequests { } impl Inserter { + /// Checks the assumptions of the logical bulk path without changing tables. + /// Unsupported requests retain ordinary insertion, including schema policy, + /// defaults, instant TTL and row-based Flow delivery. + pub async fn can_batch_metric_rows( + &self, + requests: &RowInsertRequests, + ctx: &QueryContextRef, + physical_table: &str, + ) -> Result { + if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() { + return Ok(false); + } + for request in &requests.inserts { + // The logical bulk encoder only supports scalar metric schemas. + // Check new tables too, before catalog lookup or schema changes. + if request.rows.as_ref().is_some_and(|rows| { + rows.schema.iter().any(|column| { + column.datatype_extension.is_some() + || !matches!( + ColumnDataType::try_from(column.datatype), + Ok(ColumnDataType::TimestampMillisecond + | ColumnDataType::Float64 + | ColumnDataType::String) + ) + }) + }) { + return Ok(false); + } + + let Some(table) = self + .get_table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + ) + .await? + else { + continue; + }; + let info = table.table_info(); + if info.meta.engine != METRIC_ENGINE_NAME + || info.is_ttl_instant_table() + || info + .meta + .options + .extra_options + .get(LOGICAL_TABLE_METADATA_KEY) + .map(String::as_str) + != Some(physical_table) + || info + .meta + .schema + .column_schemas() + .iter() + .any(|column| column.default_constraint().is_some()) + { + return Ok(false); + } + // Physical metric tags are nullable even when their logical schema + // is not. Arrow alignment is stricter than ordinary metric insertion. + if info + .meta + .primary_key_indices + .iter() + .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable()) + { + return Ok(false); + } + // The current Flow cache does not distinguish streaming and batch + // flows. Keep all Flow sources on the row-based delivery path. + match self.table_flownode_set_cache.get(info.table_id()).await { + Ok(None) => {} + Ok(Some(flows)) if flows.is_empty() => {} + _ => return Ok(false), + } + } + Ok(true) + } + + /// Meters an original logical-table request before bulk routing, without + /// cloning its rows or changing the request boundary used for accounting. + pub async fn meter_row_inserts( + requests: &mut RowInsertRequests, + ctx: &QueryContextRef, + ) -> Result { + let metered = InstantAndNormalInsertRequests { + normal_requests: RegionInsertRequests { + requests: requests + .inserts + .iter_mut() + .map(|request| RegionInsertRequest { + rows: request.rows.take(), + ..Default::default() + }) + .collect(), + }, + instant_requests: RegionInsertRequests::default(), + }; + let cost = write_meter!( + ctx.current_catalog(), + ctx.current_schema(), + metered, + ctx.write_rows_to_admit( + ctx.current_catalog(), + &ctx.current_schema(), + count_insert_rows(&metered)? + ), + ctx.channel() as u8 + ) + .await + .context(WriteRejectedSnafu); + for (request, region) in requests + .inserts + .iter_mut() + .zip(metered.normal_requests.requests) + { + request.rows = region.rows; + } + cost + } + pub fn new( catalog_manager: CatalogManagerRef, partition_manager: PartitionRuleManagerRef, @@ -2313,19 +2434,67 @@ mod tests { ); assert_eq!(ctx.channel(), Channel::Prometheus); } - let attempts = meter.attempts.lock().unwrap(); - let totals = attempts + { + let attempts = meter.attempts.lock().unwrap(); + let totals = attempts + .iter() + .map(|r| (r.schema.as_str(), r.rows, r.value)) + .collect::>(); + assert_eq!( + totals, + if enabled { + vec![("a", 4, 0), ("b", 2, 0)] + } else { + vec![] + } + ); + } + meter.attempts.lock().unwrap().clear(); + meter.reject.store(false, Ordering::Relaxed); + let ctx = Arc::new(QueryContext::with_channel( + CATALOG, + "logical", + Channel::Otlp, + )); + let admitted = admit_write(2, &ctx).await.unwrap(); + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric".to_string(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + }], + }; + let original = requests.clone(); + let cost = Inserter::meter_row_inserts(&mut requests, &admitted) + .await + .unwrap(); + assert_eq!(cost, if enabled { 17 } else { 0 }); + assert_eq!(requests, original); + let records = meter + .attempts + .lock() + .unwrap() .iter() - .map(|r| (r.schema.as_str(), r.rows, r.value)) + .map(|record| (record.rows, record.value)) .collect::>(); assert_eq!( - totals, + records, if enabled { - vec![("a", 4, 0), ("b", 2, 0)] + vec![(2, 0), (0, 17)] } else { vec![] } ); + meter.reject.store(true, Ordering::Relaxed); + let result = Inserter::meter_row_inserts(&mut requests, &admitted).await; + if enabled { + assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited); + } else { + assert_eq!(result.unwrap(), 0); + } + assert_eq!(requests, original); } #[test] @@ -2570,6 +2739,142 @@ mod tests { ) } + #[tokio::test] + async fn test_logical_batcher_eligibility() { + use catalog::RegisterTableRequest; + use catalog::memory::MemoryCatalogManager; + use common_meta::instruction::{CacheIdent, CreateFlow}; + use common_meta::kv_backend::KvBackendRef; + use common_meta::kv_backend::memory::MemoryKvBackend; + use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder}; + let requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "test_table".to_string(), + rows: None, + }], + }; + let original = + make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype()) + .table_info(); + for case in [ + "eligible", + "physical", + "ordinary", + "instant", + "disabled", + "hint", + "flow", + "required_tag", + "default", + ] { + let mut info = (*original).clone(); + info.meta.engine = METRIC_ENGINE_NAME.to_string(); + info.meta.options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "physical".to_string(), + ); + let mut ctx = QueryContext::arc().fork(); + match case { + "physical" => { + info.meta + .options + .extra_options + .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string()); + } + "ordinary" => info.meta.engine = "mito".to_string(), + "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant), + "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"), + "required_tag" | "default" => { + let mut columns = info.meta.schema.column_schemas().to_vec(); + let mut tag = ColumnSchema::new( + "tag", + ConcreteDataType::string_datatype(), + case != "required_tag", + ); + if case == "default" { + tag = tag + .with_default_constraint(Some(ColumnDefaultConstraint::null_value())) + .unwrap(); + } + columns.push(tag); + info.meta.schema = Arc::new( + SchemaBuilder::try_from_columns(columns) + .unwrap() + .build() + .unwrap(), + ); + info.meta.primary_key_indices = vec![2]; + } + _ => {} + } + let catalog = MemoryCatalogManager::with_default_setup(); + let table = Arc::new(table::Table::new( + Arc::new(info), + table::metadata::FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "test_table".to_string(), + table_id: 1, + table, + }) + .unwrap(); + let mut inserter = batcher_test_inserter().await; + inserter.catalog_manager = catalog; + inserter.auto_create_table = case != "disabled"; + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default()); + inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(10), + kv_backend, + )); + if case == "flow" { + inserter + .table_flownode_set_cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![1], + partition_to_peer_mapping: vec![(0, Peer::empty(1))], + })]) + .await + .unwrap(); + } + assert_eq!( + inserter + .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical") + .await + .unwrap(), + case == "eligible", + "{case}" + ); + } + } + + #[tokio::test] + async fn test_batcher_meter_preserves_request() { + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "sample".to_string(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { + values: vec![Value { + value_data: Some(api::v1::value::ValueData::F64Value(1.5)), + }], + }], + }), + }], + }; + let expected = requests.clone(); + Inserter::meter_row_inserts(&mut requests, &QueryContext::arc()) + .await + .unwrap(); + assert_eq!(requests, expected); + } + #[tokio::test] async fn test_instant_table_bypasses_batcher() { let batcher: Arc = Arc::new(UnexpectedBatcher); diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs index 10f8f558eb3..6b6ca8f7ee7 100644 --- a/src/servers/src/batcher.rs +++ b/src/servers/src/batcher.rs @@ -21,3 +21,21 @@ pub mod table; #[cfg(test)] mod test_util; + +/// Controls whether batching waits for storage before replying to the client. +const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; + +/// Returns whether pending-row batch submissions wait for the flush result +/// before replying to the client (synchronous mode), controlled by the +/// `PENDING_ROWS_BATCH_SYNC` environment variable and defaulting to `true`. +/// +/// Callers that reason about how long a remote write request may block (e.g. +/// the frontend HTTP timeout fallback) must consult this instead of +/// duplicating the env lookup. +pub fn pending_rows_batch_sync_enabled() -> bool { + std::env::var(PENDING_ROWS_BATCH_SYNC_ENV) + .ok() + .as_deref() + .and_then(|v| v.parse::().ok()) + .unwrap_or(true) +} diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index f4c764ee8a8..bd7a26da613 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -21,6 +21,7 @@ mod tables; #[cfg(test)] mod test_util; +use std::future::{Future, ready}; use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Duration; @@ -54,6 +55,7 @@ pub use crate::batcher::logical_table::region_write::{ pub use crate::batcher::logical_table::tables::{ PendingRowsSchemaAlterer, PendingRowsSchemaAltererRef, }; +use crate::batcher::pending_rows_batch_sync_enabled; use crate::error; use crate::error::{Error, Result}; use crate::metrics::{ @@ -62,24 +64,6 @@ use crate::metrics::{ const PHYSICAL_TABLE_KEY: &str = "physical_table"; -/// Whether wait for ingestion result before reply to client. -const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; - -/// Returns whether pending-row batch submissions wait for the flush result -/// before replying to the client (synchronous mode), controlled by the -/// `PENDING_ROWS_BATCH_SYNC` environment variable and defaulting to `true`. -/// -/// Callers that reason about how long a remote write request may block (e.g. -/// the frontend HTTP timeout fallback) must consult this instead of -/// duplicating the env lookup. -pub fn pending_rows_batch_sync_enabled() -> bool { - std::env::var(PENDING_ROWS_BATCH_SYNC_ENV) - .ok() - .as_deref() - .and_then(|v| v.parse::().ok()) - .unwrap_or(true) -} - const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; #[derive(Debug, Clone, Hash, Eq, PartialEq)] @@ -182,14 +166,31 @@ impl LogicalTablePendingRowsBatcher { impl LogicalTablePendingRowsBatcher { pub async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { + self.submit_with(requests, ctx, |_| ready(Ok(()))) + .await + .map(|(rows, ())| rows) + } + + /// Submits with request-level accounting after schema preparation and before + /// queue admission. Acknowledgement follows the global batching policy. + pub async fn submit_with( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + after_prepare: impl FnOnce(RowInsertRequests) -> F, + ) -> Result<(u64, T)> + where + F: Future> + Send, + { let (table_batches, total_rows) = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_build_and_align"]) .start_timer(); - self.build_and_align_table_batches(requests, &ctx).await? + self.build_and_align_table_batches(&requests, &ctx).await? }; + let prepared = after_prepare(requests).await?; if total_rows == 0 { - return Ok(0); + return Ok((0, prepared)); } // Flushes dispatch directly to datanodes, so admit once before enqueueing. @@ -269,9 +270,9 @@ impl LogicalTablePendingRowsBatcher { }; result .context(error::SubmitBatchSnafu) - .map(|()| total_rows as u64) + .map(|()| (total_rows as u64, prepared)) } else { - Ok(total_rows as u64) + Ok((total_rows as u64, prepared)) } } } diff --git a/src/servers/src/batcher/logical_table/tables.rs b/src/servers/src/batcher/logical_table/tables.rs index 4ed07f73a6e..87110366a8c 100644 --- a/src/servers/src/batcher/logical_table/tables.rs +++ b/src/servers/src/batcher/logical_table/tables.rs @@ -75,7 +75,7 @@ impl LogicalTablePendingRowsBatcher { /// RecordBatches. pub(in crate::batcher::logical_table) async fn build_and_align_table_batches( &self, - requests: RowInsertRequests, + requests: &RowInsertRequests, ctx: &QueryContextRef, ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { let catalog = ctx.current_catalog().to_string(); @@ -113,13 +113,13 @@ impl LogicalTablePendingRowsBatcher { /// Extracts non-empty `(table_name, rows)` pairs and computes total row /// count across the retained entries. pub(in crate::batcher::logical_table) fn collect_non_empty_table_rows( - requests: RowInsertRequests, - ) -> (Vec<(String, Rows)>, usize) { - let mut table_rows: Vec<(String, Rows)> = Vec::with_capacity(requests.inserts.len()); + requests: &RowInsertRequests, + ) -> (Vec<(&str, &Rows)>, usize) { + let mut table_rows: Vec<(&str, &Rows)> = Vec::with_capacity(requests.inserts.len()); let mut total_rows = 0; - for request in requests.inserts { - let Some(rows) = request.rows else { + for request in &requests.inserts { + let Some(rows) = &request.rows else { continue; }; if rows.rows.is_empty() { @@ -127,7 +127,7 @@ impl LogicalTablePendingRowsBatcher { } total_rows += rows.rows.len(); - table_rows.push((request.table_name, rows)); + table_rows.push((request.table_name.as_str(), rows)); } (table_rows, total_rows) @@ -137,15 +137,15 @@ impl LogicalTablePendingRowsBatcher { impl LogicalTablePendingRowsBatcher { /// Returns unique `(table_name, proto_schema)` pairs while keeping the /// first-seen schema for duplicate table names. - pub(in crate::batcher::logical_table) fn collect_unique_table_schemas( - table_rows: &[(String, Rows)], - ) -> Result> { + pub(in crate::batcher::logical_table) fn collect_unique_table_schemas<'a>( + table_rows: &[(&'a str, &'a Rows)], + ) -> Result> { let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); let mut seen = HashSet::new(); for (table_name, rows) in table_rows { - if seen.insert(table_name.as_str()) { - unique_tables.push((table_name.as_str(), &rows.schema)); + if seen.insert(*table_name) { + unique_tables.push((*table_name, &rows.schema)); } else { // table_rows should group rows by table name. return error::InvalidPromRemoteRequestSnafu { @@ -233,7 +233,7 @@ impl LogicalTablePendingRowsBatcher { catalog: &str, schema: &str, ctx: &QueryContextRef, - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { if plan.tables_to_create.is_empty() { @@ -298,7 +298,7 @@ impl LogicalTablePendingRowsBatcher { /// For newly created tables, re-checks all row schemas and appends alter /// operations when additional tag columns are still missing. pub(in crate::batcher::logical_table) fn enqueue_alter_for_new_tables( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { let created_tables: HashSet<&str> = plan @@ -308,11 +308,11 @@ impl LogicalTablePendingRowsBatcher { .collect(); for (table_name, rows) in table_rows { - if !created_tables.contains(table_name.as_str()) { + if !created_tables.contains(table_name) { continue; } - let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { + let Some((region_schema, _)) = plan.region_schemas.get(*table_name) else { continue; }; @@ -321,13 +321,13 @@ impl LogicalTablePendingRowsBatcher { || plan .tables_to_alter .iter() - .any(|(existing_name, _)| existing_name == table_name) + .any(|(existing_name, _)| existing_name == *table_name) { continue; } plan.tables_to_alter - .push((table_name.clone(), missing_columns)); + .push((table_name.to_string(), missing_columns)); } Ok(()) @@ -397,13 +397,13 @@ impl LogicalTablePendingRowsBatcher { /// Converts proto rows to `RecordBatch` values aligned to resolved region /// schemas and returns `(table_name, table_id, batch)` tuples. pub(in crate::batcher::logical_table) fn build_aligned_batches( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], region_schemas: &HashMap, u32)>, ) -> Result> { let mut aligned_batches = Vec::with_capacity(table_rows.len()); for (table_name, rows) in table_rows { let (region_schema, table_id) = - region_schemas.get(table_name).cloned().with_context(|| { + region_schemas.get(*table_name).cloned().with_context(|| { error::UnexpectedResultSnafu { reason: format!("Region schema not resolved for table: {}", table_name), } @@ -415,7 +415,7 @@ impl LogicalTablePendingRowsBatcher { .start_timer(); rows_to_aligned_record_batch(rows, region_schema.as_ref())? }; - aligned_batches.push((table_name.clone(), table_id, record_batch)); + aligned_batches.push((table_name.to_string(), table_id, record_batch)); } Ok(aligned_batches) @@ -531,7 +531,7 @@ mod tests { }; let (table_rows, total_rows) = - LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); + LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(&requests); assert_eq!(2, total_rows); assert_eq!(1, table_rows.len()); diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 1a46a7f10f3..67b3431a810 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -38,6 +38,7 @@ use snafu::ResultExt; use table::metadata::TableInfoRef; use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot}; +use crate::batcher::pending_rows_batch_sync_enabled; use crate::batcher::table::flow_notifier::FlowNotifier; use crate::batcher::table::metrics::PENDING_WORKERS; use crate::batcher::table::pending_batch::PendingBatch; @@ -70,6 +71,7 @@ pub struct TablePendingRowsBatcher { flush_limiter: FlushLimiter, request_limiter: RequestLimiter, worker_channel_capacity: usize, + pending_rows_batch_sync: bool, worker_idle_timeout: Duration, inserter: Arc, flow_notifier: FlowNotifier, @@ -110,6 +112,7 @@ impl TablePendingRowsBatcher { flush_limiter, request_limiter, worker_channel_capacity, + pending_rows_batch_sync: pending_rows_batch_sync_enabled(), worker_idle_timeout: flush_interval.checked_mul(3).unwrap_or(flush_interval), inserter, flow_notifier, @@ -160,8 +163,8 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { }) } - /// Waits for completed bulk writes. Cancellation does not retract an admitted - /// submission. Combined failures affect all waiters and may be partial writes. + /// Acknowledges according to the global batching policy. Cancellation does + /// not retract an admitted submission. Flush failures may be partial writes. async fn submit( &self, table_info: TableInfoRef, @@ -210,6 +213,9 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } .fail(); } + if !self.pending_rows_batch_sync { + return Ok(total_rows); + } response_rx .await .map_err(|_| { @@ -236,7 +242,7 @@ mod tests { use api::v1::region::{RegionRequest, bulk_insert_request, region_request}; use api::v1::value::ValueData; use api::v1::{ColumnDataType, Row, Rows, Value}; - use arrow::array::{Int32Array, TimestampMillisecondArray}; + use arrow::array::{ArrayRef, Int32Array, TimestampMillisecondArray}; use arrow::datatypes::Schema as ArrowSchema; use arrow::record_batch::RecordBatch; use catalog::memory::MemoryCatalogManager; @@ -320,6 +326,111 @@ mod tests { } } + #[tokio::test] + async fn test_acknowledgement_preserves_admission() { + use operator::error::UnexpectedSnafu; + + use crate::batcher::pending_rows_batch_sync_enabled; + use crate::batcher::table::batch_key_from_ctx; + use crate::batcher::table::pending_batch::notify_batches; + + for sync in [false, true] { + for fail in [false, true] { + let backend = prepare_mocked_backend().await; + let nodes = Arc::new(MockDatanodeManager::new(BulkHandler { + requests: Arc::new(Mutex::new(Vec::new())), + report_missing_row: false, + expected_skip_wal: false, + expected_schema: "public".to_string(), + })); + let inserter = Arc::new(Inserter::new( + MemoryCatalogManager::new(), + create_partition_rule_manager(backend).await, + nodes, + mock_table_flownode_cache(1, vec![]).await, + true, + )); + let mut batcher = TablePendingRowsBatcher::try_new( + Duration::from_secs(3600), + 1, + 1, + 1, + 1, + NonZeroUsize::new(1).unwrap(), + inserter, + ) + .unwrap(); + assert_eq!( + batcher.pending_rows_batch_sync, + pending_rows_batch_sync_enabled() + ); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter())); + let ctx = QueryContext::arc(); + let key = batch_key_from_ctx(&table.name, &ctx); + // Hold the worker command to control completion independently of scheduling. + let (_, receiver) = batcher.workers.get_or_create(key, 1).await; + let mut receiver = receiver.unwrap(); + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]) + .unwrap(); + let permit = batcher.acquire().await.unwrap(); + let submitter = batcher.clone(); + let submitted = + tokio::spawn(async move { submitter.submit(table, batch, ctx, permit).await }); + let WorkerCommand::Submit(pending) = + timeout(Duration::from_secs(5), receiver.recv()) + .await + .unwrap() + .unwrap(); + let mut submitted = Some(submitted); + if sync { + assert!(!submitted.as_ref().unwrap().is_finished()); + } else { + assert_eq!( + timeout(Duration::from_secs(5), submitted.take().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(), + 1 + ); + } + // Early acknowledgement must not release capacity before completion. + let acquire = batcher.acquire(); + tokio::pin!(acquire); + assert!(futures::poll!(acquire.as_mut()).is_pending()); + let result = if fail { + Err(Arc::new( + UnexpectedSnafu { + violated: "flush failed".to_string(), + } + .build(), + )) + } else { + Ok(()) + }; + notify_batches(vec![pending], result); + if let Some(submitted) = submitted { + let result = timeout(Duration::from_secs(5), submitted) + .await + .unwrap() + .unwrap(); + assert_eq!(result.is_err(), fail); + if !fail { + assert_eq!(result.unwrap(), 1); + } + } + timeout(Duration::from_secs(5), acquire) + .await + .unwrap() + .unwrap(); + } + } + } + fn rows(value: i32) -> Rows { Rows { schema: vec![ diff --git a/src/servers/src/batcher/table/batch.rs b/src/servers/src/batcher/table/batch.rs index 21ebb94a33a..e9ab339a6ac 100644 --- a/src/servers/src/batcher/table/batch.rs +++ b/src/servers/src/batcher/table/batch.rs @@ -17,6 +17,7 @@ use std::time::Instant; use arrow::compute::concat_batches; use arrow::record_batch::RecordBatch; +use common_telemetry::error; use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; @@ -55,6 +56,7 @@ pub(in crate::batcher::table) async fn flush_batch( notifier.notify(table, &combined); } Err(error) => { + error!(error; "Failed to flush table batch, rows: {}", batch.total_rows); FLUSH_FAILURES.inc(); FLUSH_DROPPED_ROWS.inc_by(batch.total_rows as u64); notify_batches(batch.submissions, Err(Arc::new(error))); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 7f2688d30cf..4d168111ade 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -215,6 +215,7 @@ pub struct HttpServer { // server configs options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, bind_addr: Option, /// What this server instance exposes. See [`HttpServerKind`]. kind: HttpServerKind, @@ -235,6 +236,9 @@ pub fn is_api_listener_path(path: &str) -> bool { is_namespace(path, HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH) || is_namespace(path, "/dashboard") } +#[derive(Clone)] +struct LogicalBatchingProtocols(Vec); + /// Sets a local-only write selector after authentication creates the context. async fn set_http_write_batching( State(protocol): State, @@ -245,8 +249,13 @@ async fn set_http_write_batching( .extensions() .get::>>() .is_some_and(|protocols| protocols.contains(&protocol)); + let logical_enabled = req + .extensions() + .get::() + .is_some_and(|protocols| protocols.0.contains(&protocol)); if let Some(ctx) = req.extensions_mut().get_mut::() { ctx.set_batching_enabled(enabled); + ctx.set_logical_batching_enabled(logical_enabled); } next.run(req).await } @@ -650,6 +659,7 @@ pub struct DashboardState { pub struct HttpServerBuilder { options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, user_provider: Option, router: Router, memory_limiter: ServerMemoryLimiter, @@ -660,13 +670,20 @@ impl HttpServerBuilder { Self { options, batching_protocols: Vec::new(), + logical_batching_protocols: Vec::new(), user_provider: None, router: Router::new(), memory_limiter: ServerMemoryLimiter::default(), } } - /// Selects HTTP write protocols allowed to use the shared batcher. + /// Selects HTTP protocols allowed to use logical-table batching. + pub fn with_logical_batching_protocols(mut self, protocols: Vec) -> Self { + self.logical_batching_protocols = protocols; + self + } + + /// Selects HTTP protocols allowed to use ordinary-table batching. pub fn with_batching_protocols(mut self, protocols: Vec) -> Self { self.batching_protocols = protocols; self @@ -894,6 +911,7 @@ impl HttpServerBuilder { HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider, shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -928,6 +946,7 @@ impl HttpServerBuilder { let internal = HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router.clone()), @@ -946,6 +965,7 @@ impl HttpServerBuilder { Some(HttpServer { options: api_options, batching_protocols: self.batching_protocols, + logical_batching_protocols: self.logical_batching_protocols, user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -1104,6 +1124,9 @@ impl HttpServer { authorize::check_http_auth, )) .layer(Extension(Arc::new(self.batching_protocols.clone()))) + .layer(Extension(LogicalBatchingProtocols( + self.logical_batching_protocols.clone(), + ))) .layer(middleware::from_fn(hints::extract_hints)) .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( @@ -2462,6 +2485,44 @@ mod batching_tests { use crate::opentsdb::codec::DataPoint; use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler}; + #[tokio::test] + async fn test_batching_selectors_are_independent() { + use axum::routing::post; + use axum::{Extension, Json, Router, middleware}; + use session::context::{QueryContext, QueryContextBuilder}; + + use crate::http::{LogicalBatchingProtocols, set_http_write_batching}; + for table in [false, true] { + for logical in [false, true] { + let app = Router::new() + .route( + "/", + post(|Extension(ctx): Extension| async move { + Json([ctx.batching_enabled(), ctx.logical_batching_enabled()]) + }), + ) + .route_layer(middleware::from_fn_with_state( + BatchingProtocol::Otlp, + set_http_write_batching, + )) + .layer(Extension(QueryContextBuilder::default().build())) + .layer(Extension(Arc::new(if table { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))) + .layer(Extension(LogicalBatchingProtocols(if logical { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))); + let client = TestClient::new(app).await; + let actual: [bool; 2] = client.post("/").send().await.json().await; + assert_eq!(actual, [table, logical]); + } + } + } + #[test] fn test_protocol_names_reject_unknown_values() { assert_eq!( diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 0bb3f0f6a31..5ebf8b36d21 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -82,6 +82,9 @@ pub struct QueryContext { /// Local-only write batching selection; never transported in protobuf extensions. #[builder(default)] batching_enabled: bool, + /// Local-only opt-in for logical metric writes, independent of ordinary-table batching. + #[builder(default)] + logical_batching_enabled: bool, /// Track which protocol the query comes from. #[builder(default)] channel: Channel, @@ -458,7 +461,17 @@ impl QueryContext { &self.configuration_parameter } - /// Whether the local HTTP entry point selected write batching. + /// Whether the local HTTP entry point selected logical-table batching. + pub fn logical_batching_enabled(&self) -> bool { + self.logical_batching_enabled + } + + /// Sets local logical-table batching selection without adding a wire-visible extension. + pub fn set_logical_batching_enabled(&mut self, enabled: bool) { + self.logical_batching_enabled = enabled; + } + + /// Whether the local HTTP entry point selected ordinary-table batching. pub fn batching_enabled(&self) -> bool { self.batching_enabled } @@ -634,6 +647,7 @@ impl QueryContextBuilder { channel, batching_enabled: self.batching_enabled.unwrap_or_default(), admitted_write: None, + logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), protocol_ctx: self.protocol_ctx.unwrap_or_default(), @@ -923,11 +937,17 @@ mod test { fn test_batching_selection_is_local_only() { let mut ctx = QueryContextBuilder::default().build(); assert!(!ctx.batching_enabled()); + assert!(!ctx.logical_batching_enabled()); + ctx.set_logical_batching_enabled(true); + assert!(ctx.clone().logical_batching_enabled()); + assert!(ctx.fork().logical_batching_enabled()); ctx.set_batching_enabled(true); assert!(ctx.clone().batching_enabled()); assert!(ctx.fork().batching_enabled()); let wire: api::v1::QueryContext = ctx.into(); - assert!(!QueryContext::from(wire).batching_enabled()); + let restored = QueryContext::from(wire); + assert!(!restored.batching_enabled()); + assert!(!restored.logical_batching_enabled()); let ctx = QueryContextBuilder::default() .set_extension("batching_enabled".to_string(), "true".to_string()) .build(); diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 80ba4cf11a5..d0de7408789 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -58,7 +58,7 @@ pub struct StandaloneOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, - /// Shared experimental ordinary-table batching; independent of Prom batching. + /// Ordinary-table batching with independent logical-table controls. pub pending_rows_batcher: PendingRowsBatcherOptions, pub jaeger: JaegerOptions, pub otlp: OtlpOptions, @@ -137,6 +137,7 @@ impl Configurable for StandaloneOptions { "wal.broker_endpoints", "event_recorder.event_types", "pending_rows_batcher.protocols", + "pending_rows_batcher.logical_table.protocols", ]) } } @@ -219,6 +220,61 @@ mod tests { use crate::options::*; + #[test] + fn test_logical_batcher_config_forwarding() { + let opts: StandaloneOptions = toml::from_str("[pending_rows_batcher]\nprotocols = ['influxdb']\n[pending_rows_batcher.logical_table]\nprotocols = ['otlp', 'prom']\npending_rows_flush_interval = '10ms'").unwrap(); + let frontend = opts.frontend_options(); + assert_eq!(frontend.pending_rows_batcher, opts.pending_rows_batcher); + assert_eq!( + frontend.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + let restored: StandaloneOptions = toml::from_str(&toml::to_string(&opts).unwrap()).unwrap(); + assert_eq!( + restored.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + assert!( + toml::from_str::( + "[pending_rows_batcher.logical_table]\nprotocols = ['logs']" + ) + .is_err() + ); + } + + #[test] + fn test_logical_batcher_protocols_from_env() { + temp_env::with_vars( + [ + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("otlp,influxdb"), + ), + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS", + Some("prom,otlp"), + ), + ], + || { + let options = + StandaloneOptions::load_layered_options(None, "STANDALONE_LOGICAL_TEST") + .unwrap(); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .protocols, + vec![ + servers::http::BatchingProtocol::Prom, + servers::http::BatchingProtocol::Otlp + ] + ); + }, + ); + } + #[test] fn test_batcher_protocols_from_env() { temp_env::with_vars( @@ -231,7 +287,7 @@ mod tests { StandaloneOptions::load_layered_options(None, "STANDALONE_BATCHER_TEST") .unwrap(); assert_eq!( - options.pending_rows_batcher.protocols, + options.pending_rows_batcher.table.protocols, vec![ servers::http::BatchingProtocol::Influxdb, servers::http::BatchingProtocol::HttpSql @@ -247,6 +303,7 @@ mod tests { assert!( !defaults .pending_rows_batcher + .table .pending_rows_batching_enabled() ); let options: StandaloneOptions = toml::from_str( @@ -259,16 +316,22 @@ flow_notification_queue_capacity = 17 "#, ) .unwrap(); - assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); - assert_eq!(options.pending_rows_batcher.protocols.len(), 2); + assert_eq!(options.pending_rows_batcher.table.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); assert_eq!( options .pending_rows_batcher + .table .flow_notification_queue_capacity .get(), 17 ); - assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + assert!( + options + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); let serialized = toml::to_string(&options).unwrap(); let parsed: StandaloneOptions = toml::from_str(&serialized).unwrap(); assert_eq!(options.influxdb, parsed.influxdb); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index dacab7ad6bd..338b03ff68f 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -503,6 +503,233 @@ WITH( Ok(()) } + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_alignment() { + use std::time::Duration; + + use common_base::Plugins; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, exponential_histogram_data_point, + }; + use prost::Message; + use servers::batcher::pending_rows_batch_sync_enabled; + use servers::http::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + // Fixed workload, only the batcher opt-in changes. Compare stored rows, + // schema evolution and visibility under the configured acknowledgement policy. + let mut results = Vec::new(); + for enabled in [false, true] { + let standalone = GreptimeDbStandaloneBuilder::new(&format!("otlp_logical_{enabled}")) + .with_logical_batcher(BatcherOptions { + protocols: if enabled { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }, + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let mut options = standalone.opts.clone(); + options.otlp.experimental_enable_exponential_histogram = true; + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_wait_flush_result"]); + let before = submissions.get_sample_count(); + for (ts, value) in [(60, 10), (120, 20)] { + let mut request = build_sum_request( + "batch.alignment", + AggregationTemporality::Cumulative, + &[(ts, value)], + ); + if ts == 120 + && let Some(metric::Data::Sum(sum)) = + &mut request.resource_metrics[0].scope_metrics[0].metrics[0].data + { + sum.data_points[0].attributes.push(keyvalue("extra", "new")); + } + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + } + if enabled { + assert_eq!( + submissions.get_sample_count() - before, + if pending_rows_batch_sync_enabled() { + 2 + } else { + 0 + }, + "logical submissions must follow the global acknowledgement policy" + ); + } + let sql = "SELECT greptime_timestamp, greptime_value, stream, extra FROM batch_alignment_total ORDER BY greptime_timestamp"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if !enabled + || pending_rows_batch_sync_enabled() + || batches.iter().map(|batch| batch.num_rows()).sum::() == 2 + { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 2 + ); + results.push(batches.pretty_print().unwrap()); + + let before_histograms = submissions.get_sample_count(); + // First create a histogram table, then reuse it in a mixed export. + // Neither request may enter the scalar-only logical batcher. + for mixed in [false, true] { + let mut request = build_sum_request( + "batch.mixed", + AggregationTemporality::Cumulative, + &[(180, 30)], + ); + let metrics = &mut request.resource_metrics[0].scope_metrics[0].metrics; + let mut histogram = metrics[0].clone(); + histogram.name = "batch.histogram".to_string(); + histogram.data = Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: AggregationTemporality::Cumulative as i32, + data_points: vec![ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: if mixed { 4_000_000_000 } else { 3_000_000_000 }, + count: 4, + sum: Some(8.0), + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![1, 2], + }), + ..Default::default() + }], + })); + if !mixed { + metrics.clear(); + } + metrics.push(histogram); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!( + response.status().as_u16(), + 200, + "enabled={enabled}, mixed={mixed}" + ); + } + assert_eq!(submissions.get_sample_count(), before_histograms); + for (sql, expected_rows) in [ + ( + "SELECT greptime_timestamp, greptime_native_histogram FROM batch_histogram ORDER BY greptime_timestamp", + 2, + ), + ( + "SELECT greptime_timestamp, greptime_value FROM batch_mixed_total ORDER BY greptime_timestamp", + 1, + ), + ] { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected_rows + ); + results.push(batches.pretty_print().unwrap()); + } + + // An existing metric attached to another physical table must use + // the original routing, not the request's default physical table. + for sql in [ + "CREATE TABLE custom_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='')", + "CREATE TABLE custom_total (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE, \"stream\" STRING PRIMARY KEY) ENGINE=metric WITH ('on_physical_table'='custom_physical')", + ] { + instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + } + instance + .metrics( + build_sum_request("custom", AggregationTemporality::Cumulative, &[(60, 7)]), + ctx.clone(), + ) + .await + .unwrap(); + let output = instance + .do_query("SELECT greptime_value FROM custom_total", ctx.clone()) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + assert!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() + .contains("7.0") + ); + + // Request-level schema policy cannot be bypassed by batching. + let mut fixed = ctx.fork(); + fixed.set_extension("auto_create_table", "false"); + assert!( + instance + .metrics( + build_sum_request( + "missing", + AggregationTemporality::Cumulative, + &[(60, 1)] + ), + Arc::new(fixed) + ) + .await + .is_err() + ); + } + assert_eq!(results[..3], results[3..]); + } + #[tokio::test(flavor = "multi_thread")] pub async fn test_otlp_on_standalone() { let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_otlp") diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 307a7f6daa2..38f546e9ccb 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -52,6 +52,7 @@ use frontend::frontend::Frontend; use frontend::instance::Instance; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; +use frontend::service_config::{BatcherOptions, PendingRowsBatcherOptions}; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; use standalone::options::StandaloneOptions; @@ -88,6 +89,7 @@ pub struct GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions, auto_create_table: bool, experimental_metric_export: bool, + logical_batcher: Option, } impl GreptimeDbStandaloneBuilder { @@ -108,6 +110,7 @@ impl GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, experimental_metric_export: false, + logical_batcher: None, } } @@ -118,6 +121,13 @@ impl GreptimeDbStandaloneBuilder { self } + /// Configures logical-table batching for integration tests. + #[must_use] + pub fn with_logical_batcher(mut self, options: BatcherOptions) -> Self { + self.logical_batcher = Some(options); + self + } + #[must_use] pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self { self.auto_create_table = auto_create_table; @@ -377,6 +387,10 @@ impl GreptimeDbStandaloneBuilder { event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, experimental_metric_export: self.experimental_metric_export, + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: self.logical_batcher.clone(), + ..Default::default() + }, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { experimental_enable_resource_info: true, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index e4e30baeba8..020e9c023c6 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2432,7 +2432,7 @@ protocols = [] pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 @@ -2453,7 +2453,7 @@ experimental_enable_prometheus_native_histogram = false pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 From e91faa9df8885506954053a2cbc9959a7e975ab6 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:52:54 +0000 Subject: [PATCH 60/63] perf(mito2): lazily decode dense primary key columns (#9226) * perf(mito2): lazily decode dense primary key columns Signed-off-by: Lei, HUANG * perf(mito2): bypass lazy decoding for full primary keys Signed-off-by: Lei, HUANG * fix(mito-codec): preserve prefix decoding errors Signed-off-by: Lei, HUANG * refactor(mito-codec): align encoded length helper naming Rename encoded_length to encoded_len and update all callers to match the other length helpers in the module. Signed-off-by: Lei, HUANG * refactor(mito2): clarify conditional dense key decoding Rename decode_dense_pk to ensure_dense_pk_decoded so callers can see that existing decoded values are preserved and only missing caches are populated. Signed-off-by: Lei, HUANG * refactor(mito-codec): share string framing in row converter Move encoded_string_len to the parent module so Dense and Sparse use the same framing helper without depending on each other. Preserve its implementation and visibility. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito-codec/src/row_converter.rs | 30 ++ src/mito-codec/src/row_converter/dense.rs | 481 ++++++++++++++---- .../src/row_converter/sparse/checked.rs | 45 +- src/mito2/benches/bench_dense_index_update.rs | 162 ++++-- src/mito2/benches/bench_pk_tag_column.rs | 117 ++++- src/mito2/src/read.rs | 127 ++++- src/mito2/src/sst/index.rs | 53 ++ src/mito2/src/sst/index/column_test.rs | 90 +++- src/mito2/src/sst/index/indexer/abort.rs | 1 + src/mito2/src/sst/index/indexer/finish.rs | 1 + src/mito2/src/sst/index/indexer/update.rs | 22 + src/mito2/src/sst/parquet/flat_format.rs | 247 ++++++++- src/mito2/src/test_util/bench_util.rs | 20 +- 13 files changed, 1167 insertions(+), 229 deletions(-) diff --git a/src/mito-codec/src/row_converter.rs b/src/mito-codec/src/row_converter.rs index fae2997182c..60bbf1935e0 100644 --- a/src/mito-codec/src/row_converter.rs +++ b/src/mito-codec/src/row_converter.rs @@ -134,6 +134,11 @@ pub trait PrimaryKeyCodec: Send + Sync + Debug { /// Returns the encoding type of the primary key. fn encoding(&self) -> PrimaryKeyEncoding; + /// Returns the dense codec for schema-aware positional access, if applicable. + fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> { + None + } + /// Decodes the primary key from the given bytes. /// /// Returns a [`CompositeValues`] that follows the primary key ordering. @@ -166,3 +171,28 @@ pub fn build_primary_key_codec_with_fields( } } } + +/// Finds the checked boundary of an Option, shared by Dense and Sparse. +/// This validates framing, not UTF-8; consumers decoding strings validate UTF-8. +pub(crate) fn encoded_string_len(bytes: &[u8]) -> memcomparable::Result { + match bytes.first().copied().ok_or(memcomparable::Error::Eof)? { + 0 => return Ok(1), + 1 => {} + marker => return Err(memcomparable::Error::InvalidTagEncoding(marker as usize)), + } + match bytes.get(1).copied().ok_or(memcomparable::Error::Eof)? { + 0 => return Ok(2), + 1 => {} + marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), + } + let mut end = 2; + loop { + let chunk = bytes.get(end..end + 9).ok_or(memcomparable::Error::Eof)?; + end += 9; + match chunk[8] { + 1..=8 => return Ok(end), + 9 => {} + marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), + } + } +} diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs index 56289a08b11..ad85f4ac8f1 100644 --- a/src/mito-codec/src/row_converter/dense.rs +++ b/src/mito-codec/src/row_converter/dense.rs @@ -27,7 +27,7 @@ use datatypes::value::ValueRef; use memcomparable::{Deserializer, Serializer}; use paste::paste; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; @@ -38,7 +38,7 @@ use crate::error::{ use crate::key_values::KeyValue; use crate::primary_key_filter::DensePrimaryKeyFilter; use crate::row_converter::{ - CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, + CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, encoded_string_len, }; /// Field to serialize and deserialize value in memcomparable format. @@ -290,84 +290,112 @@ impl SortField { bytes: &[u8], deserializer: &mut Deserializer<&[u8]>, ) -> Result { - let len = self.encoded_field_len(&bytes[deserializer.position()..])?; + let pos = deserializer.position(); + let remaining = bytes + .get(pos..) + .ok_or(memcomparable::Error::Eof) + .context(error::DeserializeFieldSnafu)?; + let len = Self::encoded_len(self.encode_data_type(), remaining)?; + if len > 1 && self.encode_data_type().is_boolean() && remaining[1] > 1 { + return Err(memcomparable::Error::InvalidBoolEncoding(remaining[1])) + .context(error::DeserializeFieldSnafu); + } deserializer.advance(len); Ok(len) } /// Checks field boundaries before any unchecked reads in memcomparable. - fn encoded_field_len(&self, bytes: &[u8]) -> Result { - let invalid = || { - error::InvalidDensePrimaryKeySnafu { - reason: "truncated field or invalid encoding", + fn encoded_len(data_type: &ConcreteDataType, bytes: &[u8]) -> Result { + let marker = bytes + .first() + .copied() + .ok_or(memcomparable::Error::Eof) + .context(error::DeserializeFieldSnafu)?; + match marker { + 0 => return Ok(1), + 1 => {} + value => { + return Err(memcomparable::Error::InvalidTagEncoding(value as usize)) + .context(error::DeserializeFieldSnafu); } - .build() - }; - match bytes.first() { - Some(0) => return Ok(1), - Some(1) => {} - _ => return Err(invalid()), } - let len = match self.encode_data_type() { - ConcreteDataType::Boolean(_) - | ConcreteDataType::Int8(_) - | ConcreteDataType::UInt8(_) => 2, + let to_skip = match data_type { + ConcreteDataType::Boolean(_) => 2, + ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2, ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3, - ConcreteDataType::Int32(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::Float32(_) - | ConcreteDataType::Date(_) => 5, - ConcreteDataType::Int64(_) - | ConcreteDataType::UInt64(_) - | ConcreteDataType::Float64(_) - | ConcreteDataType::Timestamp(_) => 9, - ConcreteDataType::Time(_) | ConcreteDataType::Duration(_) => 10, + ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5, + ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9, + ConcreteDataType::Float32(_) => 5, + ConcreteDataType::Float64(_) => 9, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) => { + // Binary is encoded as a sequence of bytes, not chunked strings. + return encoded_binary_len(bytes).context(error::DeserializeFieldSnafu); + } + ConcreteDataType::String(_) => { + return encoded_string_len(bytes).context(error::DeserializeFieldSnafu); + } + ConcreteDataType::Date(_) => 5, + ConcreteDataType::Timestamp(_) => 9, // We treat timestamp as Option + ConcreteDataType::Time(_) => 10, // i64 and 1 byte time unit + ConcreteDataType::Duration(_) => 10, ConcreteDataType::Interval(IntervalType::YearMonth(_)) => 5, ConcreteDataType::Interval(IntervalType::DayTime(_)) => 9, ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => 17, ConcreteDataType::Decimal128(_) => 19, - ConcreteDataType::Binary(_) - | ConcreteDataType::Json(_) - | ConcreteDataType::Vector(_) => { - // A serde sequence: each byte is preceded by 1, terminated by 0. - let mut pos = 1; - while bytes.get(pos) == Some(&1) { - pos += 2; - } - if bytes.get(pos) != Some(&0) { - return Err(invalid()); - } - pos + 1 - } - ConcreteDataType::String(_) => { - match bytes.get(1) { - Some(0) => return Ok(2), - Some(1) => {} - _ => return Err(invalid()), - } - let mut pos = 2; - loop { - match bytes.get(pos + 8) { - Some(1..=8) => break pos + 9, - Some(9) => pos += 9, - _ => return Err(invalid()), - } - } - } - data_type => { - return error::NotSupportedFieldSnafu { + ConcreteDataType::Null(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) => { + return NotSupportedFieldSnafu { data_type: data_type.clone(), } .fail(); } }; - snafu::ensure!( - bytes.len() >= len, - error::InvalidDensePrimaryKeySnafu { - reason: "truncated field", - } - ); - Ok(len) + if bytes.len() < to_skip { + return Err(memcomparable::Error::Eof).context(error::DeserializeFieldSnafu); + } + Ok(to_skip) + } + + /// Decodes primitive fields straight into a reference value for column + /// builders, avoiding owned Value construction and conversion. The caller has + /// already checked the encoded boundary. + fn deserialize_primitive_ref( + data_type: &ConcreteDataType, + deserializer: &mut Deserializer, + ) -> Option>> { + macro_rules! decode_primitive { + ($($variant:ident, $native:ty),* $(,)?) => { + match data_type { + $(ConcreteDataType::$variant(_) => Some( + Option::<$native>::deserialize(deserializer) + .map(ValueRef::from) + .context(error::DeserializeFieldSnafu) + ),)* + _ => None, + } + }; + } + decode_primitive!( + Boolean, bool, Int8, i8, Int16, i16, Int32, i32, Int64, i64, UInt8, u8, UInt16, u16, + UInt32, u32, UInt64, u64, Float32, f32, Float64, f64, + ) + } +} + +/// The Option marker has already been checked by SortField. +fn encoded_binary_len(bytes: &[u8]) -> memcomparable::Result { + let mut current = 1; + loop { + match bytes.get(current).copied() { + Some(0) => return Ok(current + 1), + Some(1) if bytes.get(current + 1).is_some() => current += 2, + Some(1) | None => return Err(memcomparable::Error::Eof), + Some(marker) => return Err(memcomparable::Error::InvalidSeqEncoding(marker)), + } } } @@ -444,7 +472,16 @@ impl DensePrimaryKeyCodec { return Ok(index); } let start = deserializer.position(); - let len = field.encoded_field_len(&bytes[start..])?; + // Preserve the prefix error contract used by primary-key range mapping. + let len = SortField::encoded_len(field.encode_data_type(), &bytes[start..]).map_err( + |source| match source { + error::Error::DeserializeField { .. } => error::InvalidDensePrimaryKeySnafu { + reason: "truncated field or invalid encoding", + } + .build(), + source => source, + }, + )?; // Production range mapping only needs boundaries, not allocated field values. #[cfg(any(debug_assertions, test))] field.deserialize(&mut Deserializer::new(&bytes[start..start + len]))?; @@ -470,6 +507,73 @@ impl DensePrimaryKeyCodec { Ok(values) } + /// Iterates over all source PK fields in schema order, checking each field's + /// encoded boundary before deserializing it. Stops after the first error. + /// + /// Unlike positional access, this requires no offsets or value cache. Callers + /// can consume each value immediately instead of retaining the entire key. + pub fn decode_dense_iter<'a>( + &'a self, + bytes: &'a [u8], + ) -> impl Iterator> + 'a { + self.ordered_primary_key_columns.iter().scan( + Some(Deserializer::new(bytes)), + move |state, (id, field)| { + let deserializer = state.as_mut()?; + let remaining = &bytes[deserializer.position()..]; + let decoded = SortField::encoded_len(field.encode_data_type(), remaining) + .and_then(|_| field.deserialize(deserializer).map(|value| (*id, value))); + if decoded.is_err() { + *state = None; + } + Some(decoded) + }, + ) + } + + /// Returns the column ids and sort fields in encoded primary-key order. + pub fn fields(&self) -> &[(ColumnId, SortField)] { + &self.ordered_primary_key_columns + } + + /// Decodes all fields in source order into a consumer, checking boundaries. + /// Strings borrow the reusable buffer for the duration of each callback, so + /// materializing them into column builders needs no per-value owned string. + pub fn decode_dense_with( + &self, + bytes: &[u8], + value_buf: &mut Vec, + mut consume: impl FnMut(usize, ValueRef<'_>), + ) -> Result<()> { + let mut deserializer = Deserializer::new(bytes); + for (pos, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() { + let data_type = field.encode_data_type(); + let remaining = &bytes[deserializer.position()..]; + SortField::encoded_len(data_type, remaining)?; + if data_type.is_string() && remaining[0] != 0 { + deserializer.advance(1); + deserializer + .read_bytes_into(value_buf) + .context(error::DeserializeFieldSnafu)?; + let value = std::str::from_utf8(value_buf).map_err(|err| { + error::InvalidDensePrimaryKeySnafu { + reason: format!("string is not valid UTF-8: {err}"), + } + .build() + })?; + consume(pos, ValueRef::String(value)); + } else if let Some(value) = + SortField::deserialize_primitive_ref(data_type, &mut deserializer) + { + consume(pos, value?); + } else { + let value = field.deserialize(&mut deserializer)?; + consume(pos, value.as_value_ref()); + } + } + Ok(()) + } + /// Returns the field at `pos`. /// /// # Panics @@ -488,61 +592,61 @@ impl DensePrimaryKeyCodec { offsets_buf: &mut Vec, deserializer: &mut Deserializer<&[u8]>, ) -> Result { - if pos < offsets_buf.len() { - // We computed the offset before. - let offset = offsets_buf[pos]; - deserializer.advance(offset); - return Ok(offset); - } - + ensure!( + pos < self.num_fields(), + error::InvalidDensePrimaryKeySnafu { + reason: format!( + "field position {pos} exceeds field count {}", + self.num_fields() + ), + } + ); if offsets_buf.is_empty() { - let mut offset = 0; - // Skip values before `pos`. - for i in 0..pos { - // Offset to skip before reading value i. - offsets_buf.push(offset); - let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; - offset += skip; - } - // Offset to skip before reading this value. - offsets_buf.push(offset); - Ok(offset) - } else { - // Offsets are not enough. - let value_start = offsets_buf.len() - 1; - // Advances to decode value at `value_start`. - let mut offset = offsets_buf[value_start]; - deserializer.advance(offset); - for i in value_start..pos { - // Skip value i. - let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; - // Offset for the value at i + 1. - offset += skip; - offsets_buf.push(offset); - } - Ok(offset) + // The schema bounds this cache. Reserve once even when the first + // requested field is near the end of a wide key. + offsets_buf.reserve_exact(self.num_fields()); + offsets_buf.push(0); } + // Start at the requested field if cached, otherwise resume discovery + // from the furthest known boundary. + let value_start = pos.min(offsets_buf.len() - 1); + let mut offset = offsets_buf[value_start]; + ensure!( + offset <= bytes.len(), + error::InvalidDensePrimaryKeySnafu { + reason: "cached offset exceeds key length", + } + ); + deserializer.advance(offset); + for i in value_start..pos { + offset += self.field_at(i).skip_deserialize(bytes, deserializer)?; + offsets_buf.push(offset); + } + Ok(offset) } /// Decode value at `pos` in `bytes`. /// - /// The i-th element in offsets buffer is how many bytes to skip in order to read value at `pos`. + /// The i-th element in the offsets buffer is the start of field i. The buffer + /// must only be reused for the same encoded key and codec. + #[inline] pub fn decode_value_at( &self, bytes: &[u8], pos: usize, offsets_buf: &mut Vec, ) -> Result { - let mut deserializer = Deserializer::new(bytes); - self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?; - - self.field_at(pos).deserialize(&mut deserializer) + let encoded = self.encoded_value_at(bytes, pos, offsets_buf)?; + self.field_at(pos) + .deserialize(&mut Deserializer::new(encoded)) } /// Returns the encoded bytes at `pos` in `bytes`. /// - /// The i-th element in offsets buffer is how many bytes to skip in order to read value at - /// `pos`. + /// The i-th element in the offsets buffer is the start of field i. The buffer + /// must only be reused for the same encoded key and codec. Validates the + /// framing up to this field, without validating or decoding later fields. + #[inline] pub fn encoded_value_at<'a>( &self, bytes: &'a [u8], @@ -555,6 +659,11 @@ impl DensePrimaryKeyCodec { let len = self .field_at(pos) .skip_deserialize(bytes, &mut deserializer)?; + // We already found the next field's start while validating this one. + // Reuse it instead of scanning this field again on the next lookup. + if offsets_buf.len() == pos + 1 && pos + 1 < self.num_fields() { + offsets_buf.push(offset + len); + } Ok(&bytes[offset..offset + len]) } @@ -600,6 +709,10 @@ impl PrimaryKeyCodec for DensePrimaryKeyCodec { PrimaryKeyEncoding::Dense } + fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> { + Some(self) + } + fn primary_key_filter( &self, metadata: &RegionMetadataRef, @@ -663,6 +776,78 @@ mod tests { } let decoded = encoder.decode(&result).unwrap().into_dense(); assert_eq!(decoded, row); + let mut value_buf = Vec::new(); + let mut borrowed = Vec::new(); + encoder + .decode_dense_with(&result, &mut value_buf, |_, value| { + borrowed.push(Value::from(value)) + }) + .unwrap(); + assert_eq!(borrowed, row); + let sequential = encoder + .decode_dense_iter(&result) + .collect::>>() + .unwrap(); + assert_eq!( + sequential + .iter() + .map(|(_, value)| value) + .collect::>(), + row.iter().collect::>() + ); + for end in 0..result.len() { + assert!( + encoder + .decode_dense_with(&result[..end], &mut value_buf, |_, _| {}) + .is_err() + ); + assert!( + encoder + .decode_dense_iter(&result[..end]) + .collect::>>() + .is_err(), + "truncated at {end}" + ); + } + // Every supported type must have the same boundaries in positional and + // full decoding. Exercise every truncation both at the requested field + // and while skipping it to reach a later field. + let mut offsets = Vec::new(); + let mut end = 0; + for (pos, value) in row.iter().enumerate() { + let encoded = encoder + .encoded_value_at(&result, pos, &mut offsets) + .unwrap(); + assert_eq!( + encoder + .field_at(pos) + .deserialize(&mut Deserializer::new(encoded)) + .unwrap(), + *value + ); + let start = end; + end += encoded.len(); + for len in start..end { + assert!( + encoder + .encoded_value_at(&result[..len], pos, &mut Vec::new()) + .is_err() + ); + assert!( + encoder + .decode_value_at(&result[..len], pos, &mut Vec::new()) + .is_err() + ); + if pos + 1 < row.len() { + assert!( + encoder + .encoded_value_at(&result[..len], pos + 1, &mut Vec::new()) + .is_err() + ); + } + } + } + assert_eq!(end, result.len()); let mut decoded = Vec::new(); let mut offsets = Vec::new(); // Iter two times to test offsets buffer. @@ -871,6 +1056,96 @@ mod tests { } } + #[test] + fn test_positional_decode_invalid_encoding() { + for (ty, bytes) in [ + (ConcreteDataType::int64_datatype(), vec![2]), + (ConcreteDataType::boolean_datatype(), vec![1, 2]), + (ConcreteDataType::binary_datatype(), vec![1, 2]), + (ConcreteDataType::binary_datatype(), vec![1, 1, 42, 2]), + (ConcreteDataType::string_datatype(), vec![1, 2]), + ( + ConcreteDataType::string_datatype(), + vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + ), + ( + ConcreteDataType::string_datatype(), + vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 10], + ), + ] { + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ty)), + (1, SortField::new(ConcreteDataType::int64_datatype())), + ]); + for pos in [0, 1, 2, usize::MAX] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut Vec::new()) + .is_err() + ); + assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err()); + } + assert!( + codec + .decode_dense_iter(&bytes) + .collect::>>() + .is_err() + ); + assert!( + codec + .decode_dense_with(&bytes, &mut Vec::new(), |_, _| {}) + .is_err() + ); + } + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ConcreteDataType::string_datatype())), + (1, SortField::new(ConcreteDataType::int64_datatype())), + ]); + let bytes = codec + .encode([ValueRef::String("abcdefghijk"), ValueRef::Int64(42)].into_iter()) + .unwrap(); + for pos in [2, usize::MAX] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut Vec::new()) + .is_err() + ); + assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err()); + } + for offset in [bytes.len(), bytes.len() + 1, usize::MAX] { + for pos in [0, 1] { + assert!( + codec + .encoded_value_at(&bytes, pos, &mut vec![offset]) + .is_err() + ); + assert!( + codec + .decode_value_at(&bytes, pos, &mut vec![offset]) + .is_err() + ); + } + } + let mut invalid_utf8 = bytes; + invalid_utf8[2] = 0xff; + assert!( + codec + .decode_dense_with(&invalid_utf8, &mut Vec::new(), |_, _| {}) + .is_err() + ); + assert!( + codec + .decode_dense_iter(&invalid_utf8) + .collect::>>() + .is_err() + ); + assert!( + codec + .decode_value_at(&invalid_utf8, 0, &mut Vec::new()) + .is_err() + ); + } + #[test] fn test_memcmp_dictionary() { // Test Dictionary diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs index 1e57c13cc7f..0cc9ed8211a 100644 --- a/src/mito-codec/src/row_converter/sparse/checked.rs +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -17,6 +17,7 @@ use snafu::{OptionExt, ensure}; use store_api::storage::ColumnId; use crate::error::{InvalidSparsePrimaryKeySnafu, Result}; +use crate::row_converter::encoded_string_len; use crate::row_converter::sparse::{ COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET, @@ -141,45 +142,13 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { /// Finds the end of an Option without allocating or reading its payload. /// Unlike memcomparable's skip_bytes, all advances are checked for truncated input. fn encoded_label(bytes: &[u8]) -> Result<&[u8]> { - match bytes.first() { - Some(0) => return Ok(&bytes[..1]), - Some(1) => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label null marker", - } - .fail(); + let len = encoded_string_len(bytes).map_err(|error| { + InvalidSparsePrimaryKeySnafu { + reason: error.to_string(), } - } - match bytes.get(1) { - Some(0) => return Ok(&bytes[..2]), - Some(1) => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label bytes marker", - } - .fail(); - } - } - let mut end = 2; - loop { - let chunk = bytes - .get(end..end + 9) - .context(InvalidSparsePrimaryKeySnafu { - reason: "truncated label chunk", - })?; - end += 9; - match chunk[8] { - 1..=8 => return Ok(&bytes[..end]), - 9 => {} - _ => { - return InvalidSparsePrimaryKeySnafu { - reason: "invalid label chunk length", - } - .fail(); - } - } - } + .build() + })?; + Ok(&bytes[..len]) } #[cfg(test)] diff --git a/src/mito2/benches/bench_dense_index_update.rs b/src/mito2/benches/bench_dense_index_update.rs index f81d30eb4a6..6f5ccfc080b 100644 --- a/src/mito2/benches/bench_dense_index_update.rs +++ b/src/mito2/benches/bench_dense_index_update.rs @@ -30,8 +30,8 @@ use api::v1::SemanticType; use common_test_util::temp_dir::create_temp_dir; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, StringDictionaryBuilder, TimestampMillisecondArray, - UInt8Array, UInt64Array, + ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, StringDictionaryBuilder, + TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; @@ -39,6 +39,8 @@ use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::Value; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; +use mito2::read::{Batch, BatchBuilder}; +use mito2::sst::index::Indexer; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer}; @@ -284,55 +286,115 @@ fn bench_dense_index_update(c: &mut Criterion) { } for shape in shapes { let (metadata, inverted_columns, batch) = input(&shape); - let run = |iterations| { - runtime.block_on(async { - let mut elapsed = Duration::ZERO; - for _ in 0..iterations { - let file_id = FileId::random(); - let mut inverted = shape.inverted.then(|| { - InvertedIndexer::new( - file_id, - &metadata, - intermediate.clone(), - None, - NonZeroUsize::new(SEGMENT_ROWS).unwrap(), - inverted_columns.clone(), - ) - }); - let mut bloom = if shape.bloom { - BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) - .unwrap() - } else { - None - }; - if allocations { - ALLOCATIONS.set(Some((0, 0))); + for encoded_only in [false, true] { + let name = if encoded_only { + format!("encoded/{}", shape.name) + } else { + shape.name.clone() + }; + let encoded_batches: Vec = if encoded_only { + let pk = batch + .column_by_name("__primary_key") + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + let values = pk.values().as_any().downcast_ref::().unwrap(); + (0..ROWS) + .step_by(shape.rows_per_key) + .map(|row| { + let mut builder = + BatchBuilder::new(values.value(pk.keys().value(row) as usize).to_vec()); + builder + .timestamps_array( + batch + .column_by_name("ts") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + builder + .sequences_array( + batch + .column_by_name("__sequence") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + builder + .op_types_array( + batch + .column_by_name("__op_type") + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + if shape.field { + builder + .push_field_array( + shape.tags, + batch + .column_by_name(&format!("col_{}", shape.tags)) + .unwrap() + .slice(row, shape.rows_per_key), + ) + .unwrap(); + } + builder.build().unwrap() + }) + .collect() + } else { + Vec::new() + }; + let run = |iterations| { + runtime.block_on(async { + let mut elapsed = Duration::ZERO; + for _ in 0..iterations { + let file_id = FileId::random(); + let inverted = shape.inverted.then(|| { + InvertedIndexer::new( + file_id, + &metadata, + intermediate.clone(), + None, + NonZeroUsize::new(SEGMENT_ROWS).unwrap(), + inverted_columns.clone(), + ) + }); + let bloom = if shape.bloom { + BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) + .unwrap() + } else { + None + }; + let mut indexer = Indexer::for_bench(&metadata, inverted, bloom); + // Fresh caches for each iteration; both creators share each Batch. + let mut encoded_batches = encoded_batches.clone(); + if allocations { + ALLOCATIONS.set(Some((0, 0))); + } + let start = Instant::now(); + if encoded_only { + for batch in &mut encoded_batches { + indexer.update(batch).await; + } + } else { + indexer.update_flat(&batch).await; + } + elapsed += start.elapsed(); + if allocations { + report_allocations(&name); + } + indexer.abort().await; } - let start = Instant::now(); - if let Some(indexer) = &mut inverted { - indexer.update_flat(&batch).await.unwrap(); - } - if let Some(indexer) = &mut bloom { - indexer.update_flat(&batch).await.unwrap(); - } - elapsed += start.elapsed(); - if allocations { - report_allocations(&shape.name); - } - if let Some(indexer) = &mut inverted { - indexer.abort().await.unwrap(); - } - if let Some(indexer) = &mut bloom { - indexer.abort().await.unwrap(); - } - } - elapsed - }) - }; - if allocations { - run(1); - } else { - group.bench_function(&shape.name, |b| b.iter_custom(&run)); + elapsed + }) + }; + if allocations { + run(1); + } else { + group.bench_function(&name, |b| b.iter_custom(&run)); + } } } group.finish(); diff --git a/src/mito2/benches/bench_pk_tag_column.rs b/src/mito2/benches/bench_pk_tag_column.rs index d3ded3d0b05..9e191a69511 100644 --- a/src/mito2/benches/bench_pk_tag_column.rs +++ b/src/mito2/benches/bench_pk_tag_column.rs @@ -24,6 +24,7 @@ use std::hint::black_box; use std::sync::Arc; +use std::time::Duration; use api::v1::SemanticType; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; @@ -35,9 +36,12 @@ use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; -use mito_codec::row_converter::SparsePrimaryKeyCodec; +use datatypes::value::Value; +use mito_codec::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField, SparsePrimaryKeyCodec, +}; use mito2::sst::parquet::flat_format::decode_primary_keys; -use mito2::test_util::bench_util::tag_filter_for_bench; +use mito2::test_util::bench_util::{pk_materializer_for_bench, tag_filter_for_bench}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; use store_api::storage::RegionId; @@ -237,5 +241,112 @@ fn bench_pk_tag_filters(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_pk_tag_column, bench_pk_tag_filters); +/// Measures encoded-only Dense inputs, including the cost of discovering offsets. +fn bench_dense_pk_tag_column(c: &mut Criterion) { + const TAGS: usize = 40; + let mut group = c.benchmark_group("dense_pk_tag_column"); + group.sample_size(30); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + for numeric in [false, true] { + let ty = if numeric { + ConcreteDataType::uint64_datatype() + } else { + ConcreteDataType::string_datatype() + }; + let codec = DensePrimaryKeyCodec::with_fields( + (0..TAGS) + .map(|id| (id as u32, SortField::new(ty.clone()))) + .collect(), + ); + let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1)); + for id in 0..TAGS { + metadata.push_column_metadata(ColumnMetadata { + column_id: id as u32, + column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true), + semantic_type: SemanticType::Tag, + }); + } + metadata + .push_column_metadata(ColumnMetadata { + column_id: TAGS as u32, + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + }) + .primary_key((0..TAGS as u32).collect()); + let metadata = Arc::new(metadata.build().unwrap()); + for rows_per_key in [1, 32] { + let mut keys = BinaryDictionaryBuilder::::new(); + for series in 0..ROWS / rows_per_key { + let values: Vec<_> = (0..TAGS) + .map(|id| { + if numeric { + Value::UInt64((series + id) as u64) + } else { + Value::from(format!( + "tag-{id:03}-value-{series:010}-abcdefghijklmnopqrstuvwxyz" + )) + } + }) + .collect(); + let pk = codec + .encode(values.iter().map(Value::as_value_ref)) + .unwrap(); + for _ in 0..rows_per_key { + keys.append(&pk).unwrap(); + } + } + let batch = RecordBatch::try_from_iter([ + ( + "ts", + Arc::new(TimestampMillisecondArray::from_iter_values(0..ROWS as i64)) + as ArrayRef, + ), + ("__primary_key", Arc::new(keys.finish()) as ArrayRef), + ( + "__sequence", + Arc::new(UInt64Array::from(vec![1; ROWS])) as ArrayRef, + ), + ( + "__op_type", + Arc::new(UInt8Array::from(vec![1; ROWS])) as ArrayRef, + ), + ]) + .unwrap(); + for (projection, start, count) in [ + ("first", 0, 1), + ("last", 39, 1), + ("4last", 36, 4), + ("all", 0, 40), + ] { + let kind = if numeric { "numeric" } else { "string" }; + let materialize = pk_materializer_for_bench( + metadata.clone(), + (start as u32..(start + count) as u32) + .chain([TAGS as u32]) + .collect(), + batch.schema(), + ); + assert_eq!(materialize(batch.clone()).unwrap().num_columns(), count + 4); + group.bench_function(format!("{kind}_{projection}_{rows_per_key}rpk"), |b| { + b.iter(|| { + black_box(materialize(black_box(batch.clone())).unwrap()); + }); + }); + } + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_pk_tag_column, + bench_pk_tag_filters, + bench_dense_pk_tag_column +); criterion_main!(benches); diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 72e778851ee..d5abe29f021 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -65,7 +65,7 @@ use datatypes::vectors::{ }; use futures::TryStreamExt; use futures::stream::BoxStream; -use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec}; +use mito_codec::row_converter::{CompositeValues, DensePrimaryKeyCodec, PrimaryKeyCodec}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::{ColumnId, SequenceNumber, SequenceRange}; @@ -119,6 +119,8 @@ pub struct Batch { primary_key: Vec, /// Possibly decoded `primary_key` values. Some places would decode it in advance. pk_values: Option, + /// Lazily decoded Dense columns, shared by consumers such as index creators. + dense_pk_cache: Option>, /// Timestamps of rows, should be sorted and not null. timestamps: VectorRef, /// Sequences of rows @@ -135,6 +137,13 @@ pub struct Batch { fields_idx: Option>, } +/// Only used when no fully decoded primary key was supplied by the reader. +#[derive(Debug, PartialEq, Clone)] +struct DensePkCache { + offsets: Vec, + values: Vec>, +} + impl Batch { /// Creates a new batch. pub fn new( @@ -173,12 +182,14 @@ impl Batch { /// Sets possibly decoded primary-key values. pub fn set_pk_values(&mut self, pk_values: CompositeValues) { self.pk_values = Some(pk_values); + self.dense_pk_cache = None; } /// Removes possibly decoded primary-key values. For testing only. #[cfg(any(test, feature = "test"))] pub fn remove_pk_values(&mut self) { self.pk_values = None; + self.dense_pk_cache = None; } /// Returns fields in the batch. @@ -214,6 +225,7 @@ impl Batch { Self { primary_key: vec![], pk_values: None, + dense_pk_cache: None, timestamps: Arc::new(TimestampMillisecondVectorBuilder::with_capacity(0).finish()), sequences: Arc::new(UInt64VectorBuilder::with_capacity(0).finish()), op_types: Arc::new(UInt8VectorBuilder::with_capacity(0).finish()), @@ -269,6 +281,7 @@ impl Batch { /// Be sure to update that field as well. pub fn set_primary_key(&mut self, primary_key: Vec) { self.primary_key = primary_key; + self.dense_pk_cache = None; } /// Slice the batch, returning a new batch. @@ -290,6 +303,7 @@ impl Batch { // this becomes a bottleneck. primary_key: self.primary_key.clone(), pk_values: self.pk_values.clone(), + dense_pk_cache: self.dense_pk_cache.clone(), timestamps: self.timestamps.slice(offset, length), sequences: Arc::new(self.sequences.get_slice(offset, length)), op_types: Arc::new(self.op_types.get_slice(offset, length)), @@ -772,9 +786,26 @@ impl Batch { )) } + /// Prepares a shared full-key cache when all Dense PK columns are needed. + /// Explicitly supplied decoded/defaulted values take precedence. + pub(crate) fn ensure_dense_pk_decoded(&mut self, codec: &DensePrimaryKeyCodec) -> Result<()> { + if self.pk_values.is_none() { + // A fallible iterator has no nonzero lower size hint. Reserve the + // known field count instead of growing its collected Vec per key. + let mut values = Vec::with_capacity(codec.num_fields()); + for value in codec.decode_dense_iter(&self.primary_key) { + values.push(value.context(DecodeSnafu)?); + } + self.set_pk_values(CompositeValues::Dense(values)); + } + Ok(()) + } + /// Returns the value of the column in the primary key. /// - /// Lazily decodes the primary key and caches the result. + /// Reuses predecoded values when available. Otherwise Dense keys decode only + /// the requested column, sharing offsets and values across callers; Sparse + /// keys cache the full decode. pub fn pk_col_value( &mut self, codec: &dyn PrimaryKeyCodec, @@ -782,6 +813,24 @@ impl Batch { column_id: ColumnId, ) -> Result> { if self.pk_values.is_none() { + if let Some(codec) = codec.as_dense() { + if col_idx_in_pk >= codec.num_fields() { + return Ok(None); + } + let cache = self.dense_pk_cache.get_or_insert_with(|| { + Box::new(DensePkCache { + offsets: Vec::new(), + values: vec![None; codec.num_fields()], + }) + }); + if cache.values[col_idx_in_pk].is_none() { + let value = codec + .decode_value_at(&self.primary_key, col_idx_in_pk, &mut cache.offsets) + .context(DecodeSnafu)?; + cache.values[col_idx_in_pk] = Some(value); + } + return Ok(cache.values[col_idx_in_pk].as_ref()); + } self.pk_values = Some(codec.decode(&self.primary_key).context(DecodeSnafu)?); } @@ -1098,6 +1147,7 @@ impl BatchBuilder { Ok(Batch { primary_key: self.primary_key, pk_values: None, + dense_pk_cache: None, timestamps, sequences, op_types, @@ -1254,6 +1304,79 @@ mod tests { use crate::error::Error; use crate::test_util::new_batch_builder; + #[test] + fn dense_pk_columns_are_lazy_and_reset_with_the_key() { + use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; + + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (7, SortField::new(ConcreteDataType::string_datatype())), + (3, SortField::new(ConcreteDataType::int64_datatype())), + (9, SortField::new(ConcreteDataType::string_datatype())), + ]); + let values = [ + Value::from("中文\0abcdefgh"), + Value::Int64(-42), + Value::Null, + ]; + let key = codec + .encode(values.iter().map(Value::as_value_ref)) + .unwrap(); + let mut batch = new_batch_without_fields(&[1, 2], &[1, 1], &[OpType::Put, OpType::Put]); + batch.set_primary_key(key); + for pos in [2, 0, 1, 2] { + assert_eq!( + batch.pk_col_value(&codec, pos, [7, 3, 9][pos]).unwrap(), + Some(&values[pos]) + ); + } + assert!(batch.pk_col_value(&codec, 3, 99).unwrap().is_none()); + let mut sliced = batch.slice(1, 1); + assert_eq!(sliced.pk_col_value(&codec, 0, 7).unwrap(), Some(&values[0])); + + // A corrupt unrequested suffix must not force whole-key decoding. + batch.set_primary_key(vec![0, 1]); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::Null) + ); + assert!(batch.pk_col_value(&codec, 1, 3).is_err()); + batch.set_primary_key( + codec + .encode( + [ + ValueRef::String(""), + ValueRef::Int64(8), + ValueRef::String("new"), + ] + .into_iter(), + ) + .unwrap(), + ); + assert_eq!( + batch.pk_col_value(&codec, 2, 9).unwrap(), + Some(&Value::from("new")) + ); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::from("")) + ); + + // Schema compatibility may supply already decoded/defaulted values. + batch.set_pk_values(CompositeValues::Dense(vec![(7, Value::from("default"))])); + batch.ensure_dense_pk_decoded(&codec).unwrap(); + assert_eq!( + batch.pk_col_value(&codec, 0, 7).unwrap(), + Some(&Value::from("default")) + ); + assert!(batch.pk_col_value(&codec, 1, 3).unwrap().is_none()); + batch.remove_pk_values(); + batch.ensure_dense_pk_decoded(&codec).unwrap(); + assert_eq!( + batch.pk_col_value(&codec, 1, 3).unwrap(), + Some(&Value::Int64(8)) + ); + } + fn new_batch( timestamps: &[i64], sequences: &[u64], diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index ed3ec36a473..2dd29e9803e 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -37,11 +37,13 @@ use std::sync::Arc; use bloom_filter::creator::BloomFilterIndexer; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; +use mito_codec::row_converter::DensePrimaryKeyCodec; use object_store::ObjectStore; use puffin_manager::SstPuffinManager; use smallvec::{SmallVec, smallvec}; use snafu::ResultExt; use statistics::{ByteCount, RowCount}; +use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId}; use strum::IntoStaticStr; @@ -272,6 +274,8 @@ pub struct Indexer { last_mem_fulltext_index: usize, bloom_filter_indexer: Option, last_mem_bloom_filter: usize, + /// Present only when the active creators together need every Dense PK field. + dense_pk_decoder: Option, #[cfg(feature = "vector_index")] vector_indexer: Option, #[cfg(feature = "vector_index")] @@ -280,6 +284,54 @@ pub struct Indexer { } impl Indexer { + /// Wraps real creators for update-only benchmarks without a Puffin output. + #[cfg(feature = "testing")] + pub fn for_bench( + metadata: &RegionMetadataRef, + inverted_indexer: Option, + bloom_filter_indexer: Option, + ) -> Self { + let mut indexer = Self { + region_id: metadata.region_id, + inverted_indexer, + bloom_filter_indexer, + ..Default::default() + }; + indexer.prepare_dense_pk_decoder(metadata); + indexer + } + + /// Computes demand once from active creators. Field columns and duplicate + /// requests across creators cannot turn a partial PK projection into a full one. + fn prepare_dense_pk_decoder(&mut self, metadata: &RegionMetadataRef) { + if metadata.primary_key_encoding != PrimaryKeyEncoding::Dense + || metadata.primary_key.is_empty() + { + return; + } + let requested: HashSet<_> = self + .inverted_indexer + .iter() + .flat_map(|indexer| indexer.column_ids()) + .chain( + self.bloom_filter_indexer + .iter() + .flat_map(|indexer| indexer.column_ids()), + ) + .collect(); + if metadata.primary_key.iter().all(|id| requested.contains(id)) { + self.dense_pk_decoder = Some(DensePrimaryKeyCodec::new(metadata)); + } + } + + /// Called before any creator consumes the batch, so the full decode is shared. + fn prepare_primary_key(&self, batch: &mut Batch) -> Result<()> { + if let Some(codec) = &self.dense_pk_decoder { + batch.ensure_dense_pk_decoded(codec)?; + } + Ok(()) + } + /// Updates the index with the given batch. pub async fn update(&mut self, batch: &mut Batch) { self.do_update(batch).await; @@ -398,6 +450,7 @@ impl IndexerBuilder for IndexerBuilderImpl { self.build_inverted_indexer(region_file_id.file_id(), row_group_size); indexer.fulltext_indexer = self.build_fulltext_indexer(region_file_id.file_id()).await; indexer.bloom_filter_indexer = self.build_bloom_filter_indexer(region_file_id.file_id()); + indexer.prepare_dense_pk_decoder(&self.metadata); #[cfg(feature = "vector_index")] { indexer.vector_indexer = self.build_vector_indexer(region_file_id.file_id()); diff --git a/src/mito2/src/sst/index/column_test.rs b/src/mito2/src/sst/index/column_test.rs index 146bed120c0..34ff1c19fde 100644 --- a/src/mito2/src/sst/index/column_test.rs +++ b/src/mito2/src/sst/index/column_test.rs @@ -29,7 +29,9 @@ use index::bitmap::{Bitmap, BitmapType}; use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl}; use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader}; use mito_codec::index::IndexValueCodec; -use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use mito_codec::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyCodec, PrimaryKeyCodecExt, SortField, +}; use object_store::ObjectStore; use object_store::services::Memory; use prost::Message; @@ -201,13 +203,85 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { vector_index_config: Default::default(), }; let mut expected = None; - // The legacy Batch path is an independent whole-PK decode/owned-encoding oracle. + // Demand is the union of active creators' PK columns, not their counts or + // all columns in the metadata. A malformed unneeded suffix distinguishes + // the lazy path from a premature whole-key decode. + for (inverted, bloom, disable_bloom, all_pk) in [ + (vec![0, 1], vec![], false, true), + (vec![], vec![0, 1], false, true), + (vec![0], vec![1], false, true), + (vec![0, 2], vec![0, 2], false, false), + (vec![0], vec![1], true, false), + (vec![2], vec![2], false, false), + ] { + let mut case_metadata = RegionMetadataBuilder::new(metadata.region_id); + for column in &metadata.column_metadatas { + let mut column = column.clone(); + column + .column_schema + .set_inverted_index(inverted.contains(&column.column_id)); + if !bloom.contains(&column.column_id) { + column.column_schema.unset_skipping_options().unwrap(); + } + case_metadata.push_column_metadata(column); + } + case_metadata.primary_key(metadata.primary_key.clone()); + let mut case_builder = builder.clone(); + case_builder.metadata = Arc::new(case_metadata.build().unwrap()); + if disable_bloom { + case_builder.bloom_filter_index_config.create_on_flush = crate::config::Mode::Disable; + } + let mut indexer = case_builder + .build( + RegionFileId::new(metadata.region_id, FileId::random()), + 0, + None, + ) + .await; + let make_batch = |key| { + let mut batch_builder = BatchBuilder::new(key); + batch_builder + .timestamps_array(batch.column(3).slice(0, 1)) + .unwrap(); + batch_builder + .sequences_array(batch.column(5).slice(0, 1)) + .unwrap(); + batch_builder + .op_types_array(batch.column(6).slice(0, 1)) + .unwrap(); + batch_builder.build().unwrap() + }; + let mut valid = make_batch(encoded[0].clone()); + indexer.prepare_primary_key(&mut valid).unwrap(); + if all_pk { + assert_eq!( + valid.pk_values(), + Some( + &DensePrimaryKeyCodec::new(&metadata) + .decode(&encoded[0]) + .unwrap() + ) + ); + } else { + assert!(valid.pk_values().is_none()); + } + let mut truncated = make_batch(vec![0, 1]); + assert_eq!(indexer.prepare_primary_key(&mut truncated).is_err(), all_pk); + assert!(truncated.pk_values().is_none()); + indexer.abort().await; + // Aborted creators must not retain their previous all-PK demand. + indexer.prepare_primary_key(&mut truncated).unwrap(); + } + // Explicit full decoding is the oracle for lazy encoded-only Batch inputs. // One-row slices disable run merging; a tag-only projection also exercises the no-PK fallback. - for mode in ["legacy", "runs", "single_rows", "no_pk"] { + for mode in ["eager", "legacy", "lazy", "runs", "single_rows", "no_pk"] { let file = RegionFileId::new(metadata.region_id, FileId::random()); let mut indexer = builder.build(file, 0, None).await; + if mode == "lazy" { + indexer.dense_pk_decoder = None; + } match mode { - "legacy" => { + "eager" | "legacy" | "lazy" => { for (row, key) in encoded.iter().enumerate() { let mut old = BatchBuilder::new(key.clone()); old.push_field_array(2, batch.column(2).slice(row, 1)) @@ -215,7 +289,13 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { old.timestamps_array(batch.column(3).slice(row, 1)).unwrap(); old.sequences_array(batch.column(5).slice(row, 1)).unwrap(); old.op_types_array(batch.column(6).slice(row, 1)).unwrap(); - indexer.update(&mut old.build().unwrap()).await; + let mut old = old.build().unwrap(); + if mode == "eager" { + old.set_pk_values( + DensePrimaryKeyCodec::new(&metadata).decode(key).unwrap(), + ); + } + indexer.update(&mut old).await; } } "single_rows" => { diff --git a/src/mito2/src/sst/index/indexer/abort.rs b/src/mito2/src/sst/index/indexer/abort.rs index a928eb186da..24cbb797cbd 100644 --- a/src/mito2/src/sst/index/indexer/abort.rs +++ b/src/mito2/src/sst/index/indexer/abort.rs @@ -20,6 +20,7 @@ use crate::sst::index::Indexer; impl Indexer { pub(crate) async fn do_abort(&mut self) { + self.dense_pk_decoder = None; self.do_abort_inverted_index().await; self.do_abort_fulltext_index().await; self.do_abort_bloom_filter().await; diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index c7966823ddf..823cd5c2b8a 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -27,6 +27,7 @@ use crate::sst::index::{ impl Indexer { pub(crate) async fn do_finish(&mut self) -> IndexOutput { + self.dense_pk_decoder = None; let mut output = IndexOutput::default(); let Some(mut writer) = self.build_puffin_writer().await else { diff --git a/src/mito2/src/sst/index/indexer/update.rs b/src/mito2/src/sst/index/indexer/update.rs index 095da337f49..0168bee876e 100644 --- a/src/mito2/src/sst/index/indexer/update.rs +++ b/src/mito2/src/sst/index/indexer/update.rs @@ -24,6 +24,11 @@ impl Indexer { return; } + if !self.do_prepare_primary_key(batch) { + self.do_abort().await; + return; + } + if !self.do_update_inverted_index(batch).await { self.do_abort().await; } @@ -39,6 +44,23 @@ impl Indexer { } } + /// Handles decode errors before entering asynchronous cleanup, following the + /// creators' update policy without carrying a decode error across an await. + fn do_prepare_primary_key(&self, batch: &mut Batch) -> bool { + let Err(err) = self.prepare_primary_key(batch) else { + return true; + }; + if cfg!(any(test, feature = "test")) { + panic!( + "Failed to decode primary key for indexes, region_id: {}, file_id: {}, err: {:?}", + self.region_id, self.file_id, err + ); + } else { + warn!(err; "Failed to decode primary key for indexes, region_id: {}, file_id: {}", self.region_id, self.file_id); + } + false + } + /// Returns false if the update failed. async fn do_update_inverted_index(&mut self, batch: &mut Batch) -> bool { let Some(creator) = self.inverted_indexer.as_mut() else { diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index a99fc201baf..34fb8cb873c 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -46,7 +46,7 @@ use mito_codec::row_converter::sparse::{ RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView, }; use mito_codec::row_converter::{ - CompositeValues, PrimaryKeyCodec, SparseOffsetsCache, build_primary_key_codec, + DensePrimaryKeyCodec, PrimaryKeyCodec, SparseOffsetsCache, build_primary_key_codec, }; use parquet::file::metadata::RowGroupMetaData; use snafu::{OptionExt, ResultExt, ensure}; @@ -676,8 +676,9 @@ pub(crate) fn sst_column_id_indices(metadata: &RegionMetadata) -> HashMap { - let mut decoded_pk_values = Vec::with_capacity(distinct_keys.len()); - for key in &distinct_keys { - let pk_bytes = pk_values_array.value(*key as usize); - let decoded_value = codec.decode(pk_bytes).context(DecodeSnafu)?; - decoded_pk_values.push(decoded_value); - } - DecodedKeysInner::Dense(decoded_pk_values) - } + PrimaryKeyEncoding::Dense => DecodedKeysInner::Dense { + codec: codec + .as_dense() + .context(InvalidRecordBatchSnafu { + reason: "expected dense primary key codec", + })? + .clone(), + offsets: Vec::new(), + values: pk_values_array.clone(), + distinct_keys, + }, }; Ok(DecodedPrimaryKeys { @@ -748,10 +751,16 @@ pub fn decode_primary_keys( }) } -/// Eagerly decoded dense primary keys or lazily extracted sparse primary keys. +/// Encoded primary keys and encoding-specific lookup state. enum DecodedKeysInner { - /// Dense primary keys are decoded once and shared by all tag columns. - Dense(Vec), + /// Dense keys share positional offsets across projected tag columns. + Dense { + codec: DensePrimaryKeyCodec, + values: BinaryArray, + distinct_keys: Vec, + /// Initialized by positional access; full extraction needs no offsets. + offsets: Vec>, + }, /// Sparse primary keys stay encoded; each tag column extracts only its own /// values from the raw keys. Sparse { @@ -811,7 +820,7 @@ impl DecodedPrimaryKeys { /// Gets a tag column array by column id and data type. /// /// For sparse encoding, extracts the column lazily from the encoded keys. - /// For dense encoding, uses pk_index to get values from the decoded keys. + /// For dense encoding, uses pk_index to decode only the requested field. pub fn get_tag_column( &mut self, column_id: ColumnId, @@ -827,18 +836,25 @@ impl DecodedPrimaryKeys { // Gets values from the primary key. let values_vector = match inner { - DecodedKeysInner::Dense(decoded_pk_values) => { - let mut builder = column_type.create_mutable_vector(decoded_pk_values.len()); - for decoded in decoded_pk_values { - let CompositeValues::Dense(dense) = decoded else { - return InvalidRecordBatchSnafu { - reason: "expected dense primary key values", - } - .fail(); - }; - let pk_idx = pk_index.expect("pk_index required for dense encoding"); - if pk_idx < dense.len() { - builder.push_value_ref(&dense[pk_idx].1.as_value_ref()); + DecodedKeysInner::Dense { + codec, + values, + distinct_keys, + offsets, + } => { + let pk_idx = pk_index.context(InvalidRecordBatchSnafu { + reason: "pk_index required for dense encoding", + })?; + let mut builder = column_type.create_mutable_vector(distinct_keys.len()); + if offsets.is_empty() { + *offsets = vec![Vec::new(); distinct_keys.len()]; + } + for (&key, offsets) in distinct_keys.iter().zip(offsets) { + if pk_idx < codec.num_fields() { + let value = codec + .decode_value_at(values.value(key as usize), pk_idx, offsets) + .context(DecodeSnafu)?; + builder.push_value_ref(&value.as_value_ref()); } else { builder.push_null(); } @@ -872,6 +888,49 @@ impl DecodedPrimaryKeys { } } + /// Materializes all Dense tags in source-schema order with one sequential + /// traversal per key. Values are consumed immediately by their Arrow builders. + pub fn get_dense_tag_columns(&self) -> Result> { + let DecodedKeysInner::Dense { + codec, + values, + distinct_keys, + .. + } = &self.inner + else { + return InvalidRecordBatchSnafu { + reason: "expected dense primary key values", + } + .fail(); + }; + let mut builders: Vec<_> = codec + .fields() + .iter() + .map(|(_, field)| field.data_type().create_mutable_vector(distinct_keys.len())) + .collect(); + let mut value_buf = Vec::new(); + for &key in distinct_keys { + codec + .decode_dense_with(values.value(key as usize), &mut value_buf, |pos, value| { + builders[pos].push_value_ref(&value); + }) + .context(DecodeSnafu)?; + } + codec + .fields() + .iter() + .zip(builders) + .map(|((_, field), mut builder)| { + let values = builder.to_vector().to_arrow_array(); + if field.data_type().is_string() { + Ok(Arc::new(DictionaryArray::new(self.keys_array.clone(), values)) as ArrayRef) + } else { + take(&values, &self.keys_array, None).context(ComputeArrowSnafu) + } + }) + .collect() + } + /// Gets multiple sparse tag column arrays in one pass over the distinct keys, /// sharing each key's offset discovery between the columns. /// @@ -1004,6 +1063,8 @@ impl FlatConvertFormat { }) .collect(); decoded_columns.extend(decoded_pks.get_sparse_tag_columns(&columns)?); + } else if self.projected_primary_keys.len() == self.metadata.primary_key.len() { + decoded_columns.extend(decoded_pks.get_dense_tag_columns()?); } else { for (column_id, pk_index, column_index) in &self.projected_primary_keys { let column_metadata = &self.metadata.column_metadatas[*column_index]; @@ -1086,6 +1147,138 @@ mod tests { flat_sst_arrow_schema_column_num, override_pk_field_to_binary, to_flat_sst_arrow_schema, }; + #[test] + fn dense_tag_columns_match_eager_decoding() { + use datatypes::arrow::datatypes::UInt32Type; + use datatypes::value::Value; + use datatypes::vectors::Helper; + use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; + + let columns = [ + (17, ConcreteDataType::string_datatype()), + (9, ConcreteDataType::int64_datatype()), + (2, ConcreteDataType::binary_datatype()), + ]; + let codec = DensePrimaryKeyCodec::with_fields( + columns + .iter() + .map(|(id, ty)| (*id, SortField::new(ty.clone()))) + .collect(), + ); + let rows = [ + vec![ + Value::from("中文\0abcdefgh"), + Value::Int64(-42), + Value::Binary(vec![0, 1, 255].into()), + ], + vec![Value::from(""), Value::Null, Value::Binary(vec![].into())], + vec![Value::Null, Value::Int64(i64::MAX), Value::Null], + ]; + let mut encoded: Vec<_> = rows + .iter() + .map(|row| codec.encode(row.iter().map(Value::as_value_ref)).unwrap()) + .collect(); + // Unreferenced dictionary entries must not be inspected. + encoded.push(vec![255]); + let row_keys = [2, 2, 0, 1, 0, 0]; + let pk = DictionaryArray::::new( + UInt32Array::from(row_keys.to_vec()), + Arc::new(BinaryArray::from_iter_values(&encoded)), + ); + let batch = RecordBatch::try_from_iter([ + ( + "ts", + Arc::new(TimestampMillisecondArray::from_iter_values(0..6)) as ArrayRef, + ), + ("__primary_key", Arc::new(pk) as ArrayRef), + ( + "__sequence", + Arc::new(UInt64Array::from(vec![1; 6])) as ArrayRef, + ), + ( + "__op_type", + Arc::new(UInt8Array::from(vec![1; 6])) as ArrayRef, + ), + ]) + .unwrap(); + let mut metadata = RegionMetadataBuilder::new(RegionId::new(1, 1)); + for (id, ty) in columns.iter().rev() { + metadata.push_column_metadata(ColumnMetadata { + column_id: *id, + column_schema: ColumnSchema::new(format!("tag_{id}"), ty.clone(), true), + semantic_type: SemanticType::Tag, + }); + } + metadata + .push_column_metadata(ColumnMetadata { + column_id: 23, + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + }) + .primary_key(vec![17, 9, 2]); + let metadata = Arc::new(metadata.build().unwrap()); + let format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new([17, 9, 2, 23]), + Some(batch.schema()), + "test", + false, + ) + .unwrap(); + for (offset, len) in [(0, 6), (1, 4), (3, 0)] { + let mut decoded = decode_primary_keys(&codec, &batch.slice(offset, len)).unwrap(); + let all = decoded.get_dense_tag_columns().unwrap(); + let converted = format + .convert_batch(batch.slice(offset, len), None) + .unwrap(); + // Out-of-order and repeated projections share the same key offsets. + for pos in [2, 0, 1, 0] { + let array = decoded + .get_tag_column(columns[pos].0, Some(pos), &columns[pos].1) + .unwrap(); + assert_eq!(array.to_data(), all[pos].to_data()); + assert_eq!(converted.column(pos).to_data(), all[pos].to_data()); + let actual = Helper::try_into_vector(array).unwrap(); + for (row, &key) in row_keys[offset..offset + len].iter().enumerate() { + let eager = codec + .decode_dense_without_column_id(&encoded[key as usize]) + .unwrap(); + assert_eq!(actual.get(row), eager[pos]); + } + } + // A position beyond the source schema remains NULL, never decoded + // with a newer schema's field order or type. + let missing = decoded + .get_tag_column(99, Some(3), &ConcreteDataType::int64_datatype()) + .unwrap(); + assert_eq!(missing.null_count(), len); + assert!(decoded.get_tag_column(17, None, &columns[0].1).is_err()); + } + let mut malformed_columns = batch.columns().to_vec(); + malformed_columns[1] = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 6]), + Arc::new(BinaryArray::from(vec![&[0, 1][..]])), + )); + let malformed = RecordBatch::try_new(batch.schema(), malformed_columns).unwrap(); + assert!(format.convert_batch(malformed.clone(), None).is_err()); + // A partial projection must still avoid decoding an unneeded broken suffix. + let partial = FlatReadFormat::new( + metadata, + ReadColumns::new([17, 23]), + Some(batch.schema()), + "test", + false, + ) + .unwrap(); + let partial = partial.convert_batch(malformed, None).unwrap(); + let tag = Helper::try_into_vector(partial.column(0).clone()).unwrap(); + assert!((0..6).all(|row| tag.get(row).is_null())); + } + /// Builds a `RegionMetadata` with the given number of tags and fields. fn build_metadata( num_tags: usize, diff --git a/src/mito2/src/test_util/bench_util.rs b/src/mito2/src/test_util/bench_util.rs index ac635d2799a..e0334676e5b 100644 --- a/src/mito2/src/test_util/bench_util.rs +++ b/src/mito2/src/test_util/bench_util.rs @@ -21,6 +21,7 @@ use api::v1::value::ValueData; use api::v1::{Row, Rows, SemanticType}; use datafusion_common::Column; use datafusion_expr::{Expr, lit}; +use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; @@ -30,7 +31,7 @@ use rand::seq::IndexedRandom; use store_api::metadata::{ ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef, }; -use store_api::storage::RegionId; +use store_api::storage::{ColumnId, RegionId}; use table::predicate::Predicate; use crate::memtable::KeyValues; @@ -42,6 +43,23 @@ use crate::sst::parquet::flat_format::FlatReadFormat; use crate::sst::parquet::reader::SimpleFilterContext; use crate::test_util::memtable_util::region_metadata_to_row_schema; +/// Builds the actual encoded-PK-to-flat conversion used by the SST reader. +pub fn pk_materializer_for_bench( + metadata: RegionMetadataRef, + columns: Vec, + file_schema: SchemaRef, +) -> impl Fn(RecordBatch) -> crate::error::Result { + let format = FlatReadFormat::new( + metadata, + ReadColumns::new(columns), + Some(file_schema), + "bench", + false, + ) + .unwrap(); + move |batch| format.convert_batch(batch, None) +} + /// Builds a precise-filter benchmark with tags left encoded in the primary key. pub fn tag_filter_for_bench( metadata: RegionMetadataRef, From edc81c23559e97788d01741f07bc1308b0cf4a7e Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 01:41:27 +0000 Subject: [PATCH 61/63] ci: update cargo fuzz command to use nightly toolchain explicitly (#9298) * ci: update cargo fuzz command to use nightly toolchain explicitly * ci: honor RUSTUP_TOOLCHAIN pin in fuzz orchestration script An explicit `+toolchain` argument overrides the RUSTUP_TOOLCHAIN env var in rustup precedence, so the hard-coded `cargo +nightly` in run-fuzz-targets.sh bypassed the pinned FUZZ_RUST_TOOLCHAIN (nightly-2026-03-21) configured in the workflow. - Invoke `cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run` in the script so CI uses the pinned toolchain and local runs fall back to the floating nightly - Pass RUSTUP_TOOLCHAIN through to all four fuzz-test action invocations, covering the no-prebuilt-binaries path and making the reproduce command in the summary print the exact pinned toolchain - Add test_rustup_toolchain_env_is_honored covering the pinned-env scenario for both the cargo invocation args and the summary text Addresses #9298 (review). Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .github/scripts/run-fuzz-targets-test.sh | 21 ++++++++++++++++++--- .github/scripts/run-fuzz-targets.sh | 4 ++-- .github/workflows/integration.yml | 17 +++++++++++++++-- Makefile | 8 ++++---- tests-fuzz/README.md | 6 ++++++ 5 files changed, 45 insertions(+), 11 deletions(-) diff --git a/.github/scripts/run-fuzz-targets-test.sh b/.github/scripts/run-fuzz-targets-test.sh index 5e4cda64366..9705b6bc5af 100755 --- a/.github/scripts/run-fuzz-targets-test.sh +++ b/.github/scripts/run-fuzz-targets-test.sh @@ -31,7 +31,7 @@ new_fixture() { #!/usr/bin/env bash set -euo pipefail printf '%s\t%s\t%s\n' "${GT_FUZZ_DUMP_DIR}" "$*" "${MOCK_CARGO_MARKER:-}" >>"${MOCK_CARGO_LOG}" -target="$3" +target="$4" case " ${MOCK_FAIL_TARGETS:-} " in *" ${target} "*) exit 17 ;; esac @@ -117,6 +117,20 @@ EOF set -e } +test_rustup_toolchain_env_is_honored() { + new_fixture + export RUSTUP_TOOLCHAIN=nightly-2026-03-21 + run_fixture $'fuzz_create_table' true true "fuzz_create_table" + unset RUSTUP_TOOLCHAIN + + assert_eq 1 "${fixture_status}" "pinned toolchain run status" + grep -q -- '+nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/cargo.log" || \ + fail "pinned toolchain missing from cargo invocation" + grep -q 'cargo +nightly-2026-03-21 fuzz run fuzz_create_table' "${fixture}/artifacts/summary.md" || \ + fail "pinned toolchain missing from reproduce command" + cleanup_fixture +} + test_successful_targets_run_in_order() { new_fixture run_fixture $'fuzz_create_table\nfuzz_insert' false true "" @@ -125,7 +139,7 @@ test_successful_targets_run_in_order() { assert_eq 2 "$(wc -l <"${fixture}/cargo.log" | tr -d ' ')" "cargo invocation count" assert_eq \ $'fuzz_create_table\nfuzz_insert' \ - "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^fuzz run ([^ ]+).*/\1/')" \ + "$(awk -F '\t' '{print $2}' "${fixture}/cargo.log" | sed -E 's/^\+nightly fuzz run ([^ ]+).*/\1/')" \ "target order" grep -q -- '--features=unstable' "${fixture}/cargo.log" || fail "unstable feature missing" grep -q -- '-max_total_time=120' "${fixture}/cargo.log" || fail "fuzz time missing" @@ -175,7 +189,7 @@ test_fail_fast_stops_after_first_failure() { fail "skipped target annotation missing" grep -q '### Reproduce failed targets' "${fixture}/artifacts/summary.md" || \ fail "reproduction section missing" - grep -q 'cargo fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \ + grep -q 'cargo +nightly fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \ fail "reproduction command missing" cleanup_fixture } @@ -367,5 +381,6 @@ test_collector_keeps_target_scopes_separate test_cluster_collector_honors_target_scope_and_namespace test_setup_failure_writes_artifact_contract test_setup_failure_keeps_manifest_when_collection_fails +test_rustup_toolchain_env_is_honored printf 'All fuzz orchestration script tests passed.\n' diff --git a/.github/scripts/run-fuzz-targets.sh b/.github/scripts/run-fuzz-targets.sh index 226a393e16c..205ba58c419 100755 --- a/.github/scripts/run-fuzz-targets.sh +++ b/.github/scripts/run-fuzz-targets.sh @@ -104,7 +104,7 @@ write_summary() { reproduce_args+=("GT_FUZZ_INSTANCE_ROOT_DIR=${GT_FUZZ_INSTANCE_ROOT_DIR}") fi fi - reproduce_args+=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) + reproduce_args+=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) if [[ "${fuzz_unstable}" == true ]]; then reproduce_args+=(--features=unstable) fi @@ -195,7 +195,7 @@ for ((index = 0; index < ${#targets[@]}; index++)); do fi run_args=("${fuzz_binary}" "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-artifact_prefix=${target_dir}/libfuzzer/") else - run_args=(cargo fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) + run_args=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) if [[ "${fuzz_unstable}" == true ]]; then run_args+=(--features=unstable) fi diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index c81dba199e9..4e0634fbc57 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -31,6 +31,11 @@ name: Integration CI env: CARGO_FUZZ_VERSION: "0.13.2" + # Fuzz targets require a nightly toolchain (cargo-fuzz needs + # `-Zsanitizer=fuzzer`). This pin decouples the fuzz workflow from the + # workspace default in rust-toolchain.toml, so it keeps working when the + # workspace moves to a stable toolchain. + FUZZ_RUST_TOOLCHAIN: "nightly-2026-03-21" concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -328,6 +333,7 @@ jobs: - uses: actions-rust-lang/setup-rust-toolchain@v1 with: cache: false + toolchain: ${{ env.FUZZ_RUST_TOOLCHAIN }} - name: Fuzz Dependency Cache id: fuzz-dependencies-cache uses: Swatinem/rust-cache@v2 @@ -343,7 +349,7 @@ jobs: uses: actions/cache/restore@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Report fuzz cache lookup shell: bash run: | @@ -355,6 +361,8 @@ jobs: - name: Set Rust Fuzz if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash + env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} run: | sudo apt-get install -y libfuzzer-14-dev cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked @@ -363,6 +371,7 @@ jobs: shell: bash env: CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} FUZZ_TARGETS: | fuzz_create_database fuzz_create_table @@ -410,7 +419,7 @@ jobs: uses: actions/cache/save@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', 'rust-toolchain.toml', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Pack prebuilt fuzz binaries shell: bash run: | @@ -479,6 +488,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 FUZZ_SERVICE_LOG: /tmp/greptime-fuzz-runtime/greptime.log with: @@ -537,6 +547,7 @@ jobs: - name: Run Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 GT_FUZZ_BINARY_PATH: ./bin/greptime GT_FUZZ_INSTANCE_ROOT_DIR: /tmp/greptime-fuzz-artifacts/targets/unstable_fuzz_create_table_standalone/service/instance/ @@ -820,6 +831,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 GT_KAFKA_WAL_HELPER_URL: http://127.0.0.1:8080 GT_KAFKA_ADDR: kafka.kafka-cluster.svc.cluster.local:9092 @@ -994,6 +1006,7 @@ jobs: - name: Fuzz Test uses: ./.github/actions/fuzz-test env: + RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} GT_MYSQL_ADDR: 127.0.0.1:4002 with: targets: ${{ matrix.group.targets }} diff --git a/Makefile b/Makefile index 2526f3e31eb..29801465351 100644 --- a/Makefile +++ b/Makefile @@ -212,12 +212,12 @@ sqlness-test: ## Run sqlness test. RUNS ?= 1 FUZZ_TARGET ?= fuzz_alter_table .PHONY: fuzz -fuzz: ## Run fuzz test ${FUZZ_TARGET}. - cargo fuzz run ${FUZZ_TARGET} --fuzz-dir tests-fuzz -D -s none -- -runs=${RUNS} +fuzz: ## Run fuzz test ${FUZZ_TARGET} (requires a nightly toolchain). + cargo +nightly fuzz run ${FUZZ_TARGET} --fuzz-dir tests-fuzz -D -s none -- -runs=${RUNS} .PHONY: fuzz-ls -fuzz-ls: ## List all fuzz targets. - cargo fuzz list --fuzz-dir tests-fuzz +fuzz-ls: ## List all fuzz targets (requires a nightly toolchain). + cargo +nightly fuzz list --fuzz-dir tests-fuzz .PHONY: check check: ## Cargo check all the targets. diff --git a/tests-fuzz/README.md b/tests-fuzz/README.md index cc9d7eb84ee..7945bae1583 100644 --- a/tests-fuzz/README.md +++ b/tests-fuzz/README.md @@ -6,6 +6,12 @@ cargo install cargo-fuzz ``` +Note: `cargo-fuzz` instruments targets with `-Zsanitizer=fuzzer`, so fuzz +targets must be built with a **nightly** toolchain (the workspace default +toolchain does not need to be nightly — `make fuzz` / `make fuzz-ls` +invoke `cargo +nightly` for you; CI pins its own nightly in +`FUZZ_RUST_TOOLCHAIN`). + 2. Start GreptimeDB 3. Copy the `.env.example`, which is at project root, to `.env` and change the values on need. From f3eb8e6c72e2c1583d5d4483ddc149c134966234 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 23 Sep 2026 01:52:52 +0000 Subject: [PATCH 62/63] test: cut integration test time and make the storage matrix meaningful (#9308) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cut integration test time and make the storage matrix meaningful tests-integration is ~85% of workspace test CPU, and 81% of that is the S3/S3WithCache variants of the HTTP and gRPC suites. Those suites do not touch the object store: of the 70 matrix HTTP tests only one flushed and read back an SST, so the matrix was paying real AWS round trips to re-prove protocol parsing. - Point the PR CI object-store matrix at the MinIO already started by tests-integration/fixtures. Three GT_S3_* consumers did not read GT_S3_ENDPOINT_URL and would have hit real AWS with MinIO credentials; they now do. - Add a nightly Linux job against real AWS S3, and pass GT_S3_* into the release integration-test container. The release previously ran every remote-backend case as a skip and only exercised the file backend. - Give each S3WithCache test its own read cache directory. They shared /tmp/greptimedb_cache, which the datanode wipes on startup, so a starting test deleted the read cache of a running one. - Add flush -> read-back assertions to the tests whose columns have a non-trivial SST representation: JSON/JSON2 columns, native histograms, metric-engine logical tables, and tables carrying fulltext or skipping indexes whose puffin files only exist after a flush. - Move eight tests that create no table out of the storage matrix. - Make the event recorder flush interval a constructor parameter and shorten it in the event tests, which otherwise wait a 5s window per DDL they assert on. It is skipped by serde and never reaches config files. - Drop duplicates: test_grpc_zstd_compression was a verbatim copy of test_grpc_message_size_ok and is now rewritten to assert the negotiated grpc-encoding; test_execute_copy_to_{s3,oss,gcs,azblob} were strict prefixes of their copy_from siblings; two standalone/distributed event test pairs shared one assertion body. - Fix and un-ignore stddev_by_label. stddev_pop merges partial aggregates in a parallelism-dependent order, so its last digits are unstable; the test now compares values with a tolerance. - Rebase the jaeger v1 fixture on the current instant. It carries ttl=7d with 2025 timestamps, so its rows were only readable as long as they stayed in the memtable. Signed-off-by: Dennis Zhuang * test: address review — wire nightly real-S3 job into check-status, keep the short event interval The nightly `check-status` job did not depend on the new real-S3 job, so a failure there would not have reached the status or Slack notification. In database_ddl_event the short interval was set by a first `with_event_recorder_options` call and then overwritten by the pre-existing one, which carries `..Default::default()`. Merged into a single call. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../actions/build-linux-artifacts/action.yml | 23 + .github/workflows/dev-build.yml | 8 + .github/workflows/nightly-build.yml | 8 + .github/workflows/nightly-ci.yml | 53 ++- .github/workflows/release.yml | 8 + .github/workflows/rust.yml | 34 +- Makefile | 1 + src/common/event-recorder/src/recorder.rs | 30 +- src/frontend/src/instance/builder.rs | 1 + src/meta-srv/src/metasrv/builder.rs | 1 + src/object-store/src/test_util.rs | 29 -- src/object-store/tests/object_store_test.rs | 10 +- tests-integration/src/cluster.rs | 14 + tests-integration/src/test_util.rs | 79 ++-- tests-integration/src/tests/instance_test.rs | 198 +-------- tests-integration/src/tests/promql_test.rs | 49 ++- .../tests/admin_function_event.rs | 25 +- tests-integration/tests/database_ddl_event.rs | 7 +- .../tests/export_logical_tables.rs | 9 +- tests-integration/tests/flow_ddl_event.rs | 5 +- tests-integration/tests/gc_event.rs | 3 + tests-integration/tests/grpc.rs | 147 ++++++- tests-integration/tests/http.rs | 396 ++++++++++++++++-- tests-integration/tests/main.rs | 11 + .../tests/reconciliation_event.rs | 4 + tests-integration/tests/repartition_event.rs | 2 + tests-integration/tests/table_ddl_event.rs | 5 +- tests-integration/tests/view_ddl_event.rs | 16 +- tests-integration/tests/wal_prune_event.rs | 2 + 29 files changed, 805 insertions(+), 373 deletions(-) diff --git a/.github/actions/build-linux-artifacts/action.yml b/.github/actions/build-linux-artifacts/action.yml index 3cb9c439551..c5212c28998 100644 --- a/.github/actions/build-linux-artifacts/action.yml +++ b/.github/actions/build-linux-artifacts/action.yml @@ -31,6 +31,22 @@ inputs: description: Build GreptimeDB with large page size (65536). required: false default: 'false' + aws-ci-test-bucket: + description: Bucket used by the object-store integration tests + required: false + default: '' + aws-region: + description: Region of the test bucket + required: false + default: '' + aws-access-key-id: + description: Access key id for the test bucket + required: false + default: '' + aws-secret-access-key: + description: Secret access key for the test bucket + required: false + default: '' runs: using: composite steps: @@ -38,6 +54,13 @@ runs: if: ${{ inputs.disable-run-tests == 'false' }} shell: bash # NOTE: If the BUILD_JOBS > 4, it's always OOM in EC2 instance. + # Without GT_S3_*, every remote object-store case is skipped and the + # release only ever exercises the local file backend. + env: + GT_S3_BUCKET: ${{ inputs.aws-ci-test-bucket }} + GT_S3_REGION: ${{ inputs.aws-region }} + GT_S3_ACCESS_KEY_ID: ${{ inputs.aws-access-key-id }} + GT_S3_ACCESS_KEY: ${{ inputs.aws-secret-access-key }} run: | cd ${{ inputs.working-dir }} && \ make run-it-in-container BUILD_JOBS=4 \ diff --git a/.github/workflows/dev-build.yml b/.github/workflows/dev-build.yml index d03fbeff14a..563435b708e 100644 --- a/.github/workflows/dev-build.yml +++ b/.github/workflows/dev-build.yml @@ -183,6 +183,10 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: ${{ inputs.large-page-size }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -217,6 +221,10 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: ${{ inputs.large-page-size }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml index 14ebb6e715e..c78ca4dc5b9 100644 --- a/.github/workflows/nightly-build.yml +++ b/.github/workflows/nightly-build.yml @@ -151,6 +151,10 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -173,6 +177,10 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index d4afa2778b8..8441ed02312 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -71,6 +71,57 @@ jobs: path: C:\Users\RUNNER~1\AppData\Local\Temp\sqlness* retention-days: 3 + # PR CI runs the object-store matrix against a local MinIO. This job is where + # the same matrix meets real AWS S3. + test-on-linux-real-s3: + name: Run tests on Linux against real S3 + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + runs-on: ubuntu-22.04-arm + timeout-minutes: 90 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - uses: rui314/setup-mold@v1 + - name: Install toolchain + uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "coverage-test" + cache-all-crates: "true" + save-if: false + - name: Install latest nextest release + uses: taiki-e/install-action@nextest + - name: Setup external services + working-directory: tests-integration/fixtures + run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait + - name: Run nextest cases + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions + env: + CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" + RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB + CARGO_INCREMENTAL: 0 + GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} + GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 + GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 + GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres + GT_POSTGRES15_ENDPOINTS: postgres://test_user:test_password@127.0.0.1:5433/postgres + GT_POSTGRES15_SCHEMA: test_schema + GT_MYSQL_ENDPOINTS: mysql://greptimedb:admin@127.0.0.1:3306/mysql + GT_KAFKA_ENDPOINTS: 127.0.0.1:9092 + GT_KAFKA_SASL_ENDPOINTS: 127.0.0.1:9093 + UNITTEST_LOG_DIR: "__unittest_logs" + test-on-windows: name: Run tests on Windows if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} @@ -130,7 +181,7 @@ jobs: check-status: name: Check status - needs: [sqlness-test, sqlness-windows, test-on-windows] + needs: [sqlness-test, sqlness-windows, test-on-windows, test-on-linux-real-s3] if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest outputs: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index acf52e241c2..6c0c36c895e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -204,6 +204,10 @@ jobs: disable-run-tests: ${{ env.DISABLE_RUN_TESTS }} image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -227,6 +231,10 @@ jobs: image-registry: ${{ vars.ECR_IMAGE_REGISTRY }} image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: true + aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} + aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1a093a2d2a5..e4483b4a42b 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -206,15 +206,14 @@ jobs: RUST_BACKTRACE: 1 RUST_MIN_STACK: 8388608 # 8MB CARGO_INCREMENTAL: 0 - GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} - GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - GT_MINIO_BUCKET: greptime - GT_MINIO_ACCESS_KEY_ID: superpower_ci_user - GT_MINIO_ACCESS_KEY: superpower_password - GT_MINIO_REGION: us-west-2 - GT_MINIO_ENDPOINT_URL: http://127.0.0.1:9000 + # Runs the object-store matrix against the local MinIO from + # tests-integration/fixtures. Real AWS S3 is covered by the nightly + # jobs and by the release build. + GT_S3_BUCKET: greptime + GT_S3_ACCESS_KEY_ID: superpower_ci_user + GT_S3_ACCESS_KEY: superpower_password + GT_S3_REGION: us-west-2 + GT_S3_ENDPOINT_URL: http://127.0.0.1:9000 GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres @@ -263,15 +262,14 @@ jobs: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 CARGO_INCREMENTAL: 0 - GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} - GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - GT_MINIO_BUCKET: greptime - GT_MINIO_ACCESS_KEY_ID: superpower_ci_user - GT_MINIO_ACCESS_KEY: superpower_password - GT_MINIO_REGION: us-west-2 - GT_MINIO_ENDPOINT_URL: http://127.0.0.1:9000 + # Runs the object-store matrix against the local MinIO from + # tests-integration/fixtures. Real AWS S3 is covered by the nightly + # jobs and by the release build. + GT_S3_BUCKET: greptime + GT_S3_ACCESS_KEY_ID: superpower_ci_user + GT_S3_ACCESS_KEY: superpower_password + GT_S3_REGION: us-west-2 + GT_S3_ENDPOINT_URL: http://127.0.0.1:9000 GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 GT_POSTGRES_ENDPOINTS: postgres://greptimedb:admin@127.0.0.1:5432/postgres diff --git a/Makefile b/Makefile index 29801465351..7176eafd0bb 100644 --- a/Makefile +++ b/Makefile @@ -264,6 +264,7 @@ stop-etcd: ## Stop single node etcd for testing purpose. run-it-in-container: start-etcd ## Run integration tests in dev-builder. docker run --network=host \ -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git -v /tmp:/tmp \ + -e GT_S3_BUCKET -e GT_S3_ACCESS_KEY_ID -e GT_S3_ACCESS_KEY -e GT_S3_REGION \ -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} \ make test sqlness-test BUILD_JOBS=${BUILD_JOBS} diff --git a/src/common/event-recorder/src/recorder.rs b/src/common/event-recorder/src/recorder.rs index b6c18ed08de..3ea5c7e55a8 100644 --- a/src/common/event-recorder/src/recorder.rs +++ b/src/common/event-recorder/src/recorder.rs @@ -291,17 +291,29 @@ pub struct EventRecorderOptions { skip_serializing_if = "event_type_filter_is_all" )] pub event_types: EventTypeFilterRef, + /// How long buffered events may wait before reaching the event handler. + /// + /// Not a configuration option: it is skipped by serde so it never appears in + /// config files or `config/config.md`. Integration tests shorten it so they + /// do not pay a flush window per DDL they assert on. + #[serde(skip, default = "default_flush_interval")] + pub flush_interval: Duration, } fn default_events_table_ttl() -> Duration { DEFAULT_EVENTS_TABLE_TTL } +fn default_flush_interval() -> Duration { + DEFAULT_FLUSH_INTERVAL_SECONDS +} + impl Default for EventRecorderOptions { fn default() -> Self { Self { ttl: DEFAULT_EVENTS_TABLE_TTL, event_types: Arc::new(EventTypeFilter::All), + flush_interval: DEFAULT_FLUSH_INTERVAL_SECONDS, } } } @@ -321,13 +333,18 @@ pub struct EventRecorderImpl { impl EventRecorderImpl { pub fn new(event_handler: Box) -> Self { - Self::with_event_type_filter(event_handler, Arc::new(EventTypeFilter::All)) + Self::with_event_type_filter( + event_handler, + Arc::new(EventTypeFilter::All), + DEFAULT_FLUSH_INTERVAL_SECONDS, + ) } /// Creates an event recorder with an event-type filter. pub fn with_event_type_filter( event_handler: Box, event_types: EventTypeFilterRef, + flush_interval: Duration, ) -> Self { let (tx, rx) = channel(DEFAULT_CHANNEL_SIZE); let cancel_token = CancellationToken::new(); @@ -339,13 +356,9 @@ impl EventRecorderImpl { cancel_token: cancel_token.clone(), }; - let processor = EventProcessor::new( - rx, - event_handler, - DEFAULT_FLUSH_INTERVAL_SECONDS, - DEFAULT_MAX_RETRY_TIMES, - ) - .with_cancel_token(cancel_token); + let processor = + EventProcessor::new(rx, event_handler, flush_interval, DEFAULT_MAX_RETRY_TIMES) + .with_cancel_token(cancel_token); // Spawn a background task to process the events. let handle = tokio::spawn(async move { @@ -611,6 +624,7 @@ mod tests { count: count.clone(), }), event_type_filter.clone(), + DEFAULT_FLUSH_INTERVAL_SECONDS, ); assert!(Arc::ptr_eq( diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index a3cfbcdaffe..d9b86ac5a05 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -381,6 +381,7 @@ impl FrontendBuilder { self.options.event_recorder.ttl, )), self.options.event_recorder.event_types.clone(), + self.options.event_recorder.flush_interval, )); admin_event_recorder.install(&event_recorder); diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index b40f06e3cfe..cbc859716a4 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -250,6 +250,7 @@ impl MetasrvBuilder { let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter( Box::new(EventHandlerImpl::new(event_inserter)), options.event_recorder.event_types.clone(), + options.event_recorder.flush_interval, )); let selector = selector.unwrap_or_else(|| Arc::new(LeaseBasedSelector)); diff --git a/src/object-store/src/test_util.rs b/src/object-store/src/test_util.rs index af68978b0e4..936ecba5bac 100644 --- a/src/object-store/src/test_util.rs +++ b/src/object-store/src/test_util.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::env; - use crate::{ObjectStore, Result}; /// Temp folder for object store test @@ -35,30 +33,3 @@ impl TempFolder { self.store.delete_with(&self.path).recursive(true).await } } - -/// Test s3 config from environment variables -#[derive(Debug)] -pub struct TestS3Config { - pub root: String, - pub access_key_id: String, - pub secret_access_key: String, - pub bucket: String, - pub region: Option, -} - -/// Returns s3 test config, return None if not found. -pub fn s3_test_config() -> Option { - if let Ok(b) = env::var("GT_S3_BUCKET") - && !b.is_empty() - { - return Some(TestS3Config { - root: uuid::Uuid::new_v4().to_string(), - access_key_id: env::var("GT_S3_ACCESS_KEY_ID").ok()?, - secret_access_key: env::var("GT_S3_ACCESS_KEY").ok()?, - bucket: env::var("GT_S3_BUCKET").ok()?, - region: Some(env::var("GT_S3_REGION").ok()?), - }); - } - - None -} diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index 9cf03fe42c5..735d2c56f15 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -284,13 +284,21 @@ async fn test_s3_backend() -> Result<()> { let root = uuid::Uuid::new_v4().to_string(); - let builder = S3::default() + let mut builder = S3::default() .root(&root) .access_key_id(&env::var("GT_S3_ACCESS_KEY_ID")?) .secret_access_key(&env::var("GT_S3_ACCESS_KEY")?) .region(&env::var("GT_S3_REGION")?) .bucket(&bucket); + // Honors an S3-compatible endpoint (MinIO in CI) so the test does not + // fall back to resolving the bucket against real AWS. + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") + && !endpoint.is_empty() + { + builder = builder.endpoint(&endpoint); + } + let store = ObjectStore::new(builder).unwrap(); let store = object_store::util::with_instrument_layers(store, false); diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 3e53b6104a0..bc2ed6f20a4 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -36,6 +36,7 @@ use client::client_manager::NodeClients; use cmd::frontend::create_heartbeat_task; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; @@ -174,6 +175,7 @@ pub struct GreptimeDbClusterBuilder { shared_home_dir: Option>, meta_selector: Option, local_file_access: LocalFileAccess, + event_recorder_options: EventRecorderOptions, } impl GreptimeDbClusterBuilder { @@ -209,6 +211,7 @@ impl GreptimeDbClusterBuilder { shared_home_dir: None, meta_selector: None, local_file_access: LocalFileAccess::default(), + event_recorder_options: EventRecorderOptions::default(), } } @@ -253,6 +256,15 @@ impl GreptimeDbClusterBuilder { self } + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + /// Sets whether the test frontend automatically creates tables on write. #[must_use] pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { @@ -303,6 +315,7 @@ impl GreptimeDbClusterBuilder { ..Default::default() }, gc: self.metasrv_gc_config.clone(), + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; @@ -550,6 +563,7 @@ impl GreptimeDbClusterBuilder { fn build_frontend_options(&self) -> FrontendOptions { let mut fe_opts = FrontendOptions { auto_create_table: self.frontend_auto_create_table, + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 562bd48cdb0..fd02db2e5a1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -31,6 +31,7 @@ use client::{Client, Database}; use common_base::Plugins; use common_catalog::consts::MIN_USER_TABLE_ID; use common_config::Configurable; +use common_event_recorder::EventRecorderOptions; #[cfg(test)] use common_meta::DatanodeId; use common_meta::key::TableMetadataManager; @@ -328,6 +329,18 @@ impl StorageType { } } +/// Event recorder options for tests. +/// +/// The production flush interval is 5s, and the event tests interleave "run a +/// DDL, wait for its event, run the next DDL", so each barrier costs a full +/// window. A short interval removes that wait without changing what is asserted. +pub fn test_event_recorder_options() -> EventRecorderOptions { + EventRecorderOptions { + flush_interval: Duration::from_millis(100), + ..Default::default() + } +} + fn s3_test_config() -> S3Config { S3Config { connection: S3Connection { @@ -363,7 +376,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Gcs::from(&gcs_config.connection); let config = ObjectStoreConfig::Gcs(gcs_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Gcs(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Azblob => { let azblob_config = AzblobConfig { @@ -381,7 +394,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Azblob::from(&azblob_config.connection); let config = ObjectStoreConfig::Azblob(azblob_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Azblob(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Oss => { let oss_config = OssConfig { @@ -398,32 +411,57 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Oss::from(&oss_config.connection); let config = ObjectStoreConfig::Oss(oss_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Oss(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::S3 | StorageType::S3WithCache => { let mut s3_config = s3_test_config(); - if *store_type == StorageType::S3WithCache { - s3_config.cache.cache_path = "/tmp/greptimedb_cache".to_string(); + // The datanode wipes `/cache/object/read` on startup, so a + // path shared between concurrently running tests lets a starting test + // delete the read cache of a running one. + let cache_dir = if *store_type == StorageType::S3WithCache { + let dir = create_temp_dir("gt_s3_read_cache"); + s3_config.cache.cache_path = dir.path().to_string_lossy().to_string(); + Some(dir) } else { s3_config.cache.enable_read_cache = false; - } + None + }; let builder = S3::from(&s3_config.connection); let config = ObjectStoreConfig::S3(s3_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::S3(TempFolder::new(&store, "/"))) + ( + config, + TempDirGuard { + remote: Some(TempFolder::new(&store, "/")), + local_cache: cache_dir, + }, + ) } - StorageType::File => (ObjectStoreConfig::File(FileConfig {}), TempDirGuard::None), + StorageType::File => ( + ObjectStoreConfig::File(FileConfig {}), + TempDirGuard::default(), + ), } } -pub enum TempDirGuard { - None, - S3(TempFolder), - Oss(TempFolder), - Azblob(TempFolder), - Gcs(TempFolder), +#[derive(Default)] +pub struct TempDirGuard { + /// Prefix to wipe from the remote object store, absent for the file backend. + remote: Option, + /// Local read cache directory. Only held so it is removed when the guard drops. + #[allow(dead_code)] + local_cache: Option, +} + +impl TempDirGuard { + fn remote(folder: TempFolder) -> Self { + Self { + remote: Some(folder), + local_cache: None, + } + } } pub struct TestGuard { @@ -446,11 +484,7 @@ pub struct StorageGuard(pub TempDirGuard); impl TestGuard { pub async fn remove_all(&mut self) { for storage_guard in self.storage_guards.iter_mut() { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = &mut storage_guard.0 - { + if let Some(guard) = &mut storage_guard.0.remote { guard.remove_all().await.unwrap() } } @@ -465,11 +499,8 @@ impl Drop for TestGuard { common_runtime::spawn_global(async move { let mut errors = vec![]; for guard in guards { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = guard.0 - && let Err(e) = guard.remove_all().await + if let Some(remote) = guard.0.remote + && let Err(e) = remote.remove_all().await { errors.push(e); } diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index dc4d3898860..caac6d33238 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -2195,47 +2195,6 @@ async fn test_delete(instance: Arc) { check_output_stream(output, expect).await; } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_s3(instance: Arc) { - if let Ok(bucket) = env::var("GT_S3_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - // setups - assert!(matches!(execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await.data, OutputData::AffectedRows(0))); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - let key_id = env::var("GT_S3_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_S3_ACCESS_KEY").unwrap(); - let region = env::var("GT_S3_REGION").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - // exports - let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_s3(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2268,9 +2227,21 @@ async fn test_execute_copy_from_s3(instance: Arc) { let key = env::var("GT_S3_ACCESS_KEY").unwrap(); let region = env::var("GT_S3_REGION").unwrap(); + let mut connection = format!( + "ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}'", + key_id, key, region + ); + // Honors an S3-compatible endpoint (MinIO in CI) so the test does not + // fall back to resolving the bucket against real AWS. + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") + && !endpoint.is_empty() + { + connection = format!("{},ENDPOINT='{}'", connection, endpoint); + } + let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region + "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION ({})", + bucket, root, connection ); let output = execute_sql(&instance, ©_to_stmt).await.data; @@ -2315,10 +2286,7 @@ async fn test_execute_copy_from_s3(instance: Arc) { .data, OutputData::AffectedRows(0) )); - let sql = format!( - "{} CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - test.sql, key_id, key, region, - ); + let sql = format!("{} CONNECTION ({})", test.sql, connection); let output = execute_sql(&instance, &sql).await.data; assert!(matches!(output, OutputData::AffectedRows(2))); @@ -2340,50 +2308,6 @@ async fn test_execute_copy_from_s3(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_oss(instance: Arc) { - if let Ok(bucket) = env::var("GT_OSS_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let key_id = env::var("GT_OSS_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_OSS_ACCESS_KEY").unwrap(); - let endpoint = env::var("GT_OSS_ENDPOINT").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'oss://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", - bucket, root, key_id, key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_oss(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2490,54 +2414,6 @@ async fn test_execute_copy_from_oss(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_gcs(instance: Arc) { - if let (Ok(bucket), Ok(scope), Ok(credential)) = ( - env::var("GT_GCS_BUCKET"), - env::var("GT_GCS_SCOPE"), - env::var("GT_GCS_CREDENTIAL"), - ) && !bucket.is_empty() - { - let endpoint = env::var("GT_GCS_ENDPOINT").unwrap_or_default(); - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - let mut connection = format!("SCOPE='{}',CREDENTIAL='{}'", scope, credential); - if !endpoint.is_empty() { - connection = format!("{},ENDPOINT='{}'", connection, endpoint); - } - - let copy_to_stmt = format!( - "Copy demo TO 'gcs://{}/{}/export/demo.parquet' CONNECTION ({})", - bucket, root, connection - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_gcs(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2646,50 +2522,6 @@ async fn test_execute_copy_from_gcs(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_azblob(instance: Arc) { - if let (Ok(container), Ok(account_name), Ok(account_key), Ok(endpoint)) = ( - env::var("GT_AZBLOB_CONTAINER"), - env::var("GT_AZBLOB_ACCOUNT_NAME"), - env::var("GT_AZBLOB_ACCOUNT_KEY"), - env::var("GT_AZBLOB_ENDPOINT"), - ) && !container.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'azblob://{}/{}/export/demo.parquet' CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", - container, root, account_name, account_key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_azblob(instance: Arc) { common_telemetry::init_default_ut_logging(); diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index c816583f135..74114bbd6c4 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -644,28 +644,57 @@ async fn two_aggregators_combined_aggrs(instance: Arc) { // eval instant at 50m stddev by (instance)(http_requests) // {instance="0"} 223.60679774998 // {instance="1"} 223.60679774998 +// +// The values are compared with a tolerance rather than against a rendered table: +// `stddev_pop` merges partial aggregates in a parallelism-dependent order, so the +// last digits of its output are not stable across runs. #[apply(both_instances_cases)] -#[ignore = "TODO(ruihang): fix this case"] async fn stddev_by_label(instance: Arc) { let instance = instance.frontend(); - create_insert_query_assert( + execute_all(&instance, AGGREGATORS_CREATE_TABLE, QueryContext::arc()).await; + execute_all(&instance, AGGREGATORS_INSERT_DATA, QueryContext::arc()).await; + + let batches = promql_query_as_batches( instance, - AGGREGATORS_CREATE_TABLE, - AGGREGATORS_INSERT_DATA, r#"stddev by (instance)(http_requests)"#, + None, + QueryContext::arc(), UNIX_EPOCH, unix_epoch_plus_100s(), Duration::from_secs(60), Duration::from_secs(0), - "+----------+---------------------+--------------------------------+\ - \n| instance | ts | STDDEVPOP(http_requests.value) |\ - \n+----------+---------------------+--------------------------------+\ - \n| 0 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n| 1 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n+----------+---------------------+--------------------------------+", ) .await; + + let batches = batches.iter().collect::>(); + assert_eq!(batches.len(), 1); + // Tag columns come back dictionary-encoded. + let instances = cast( + batches[0].column_by_name("instance").unwrap(), + &DataType::Utf8, + ) + .unwrap(); + let instances = instances.as_any().downcast_ref::().unwrap(); + let values = batches[0] + .column_by_name("stddev_pop(http_requests.value)") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + let mut actual = (0..values.len()) + .map(|i| (instances.value(i), values.value(i))) + .collect::>(); + actual.sort_by(|a, b| a.0.cmp(b.0)); + assert_eq!(actual.len(), 2, "{}", batches[0].pretty_print()); + for (expected_instance, (instance, value)) in ["0", "1"].into_iter().zip(actual) { + assert_eq!(instance, expected_instance); + assert!( + (value - 223.606_797_749_979).abs() < 1e-9, + "instance {instance}: {value}" + ); + } } // This is not derived from prometheus diff --git a/tests-integration/tests/admin_function_event.rs b/tests-integration/tests/admin_function_event.rs index d669d2900ad..43f7445ca2f 100644 --- a/tests-integration/tests/admin_function_event.rs +++ b/tests-integration/tests/admin_function_event.rs @@ -14,7 +14,6 @@ use std::sync::Arc; -use common_test_util::temp_dir::create_temp_dir; use frontend::instance::Instance; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; @@ -23,9 +22,8 @@ use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::admin::Admin; use sql::statements::statement::Statement; -use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::assert_eventually_eq; @@ -35,31 +33,12 @@ const TABLE_NAME: &str = "admin_function_event_test"; async fn test_standalone_admin_function_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_admin_function_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; assert_admin_function_events(standalone.fe_instance()).await; } -#[tokio::test(flavor = "multi_thread")] -async fn test_distributed_admin_function_events() { - let store_type = StorageType::File; - if !store_type.test_on() { - return; - } - - common_telemetry::init_default_ut_logging(); - let (store_config, _guard) = get_test_store_config(&store_type); - let home_dir = create_temp_dir("test_distributed_admin_function_events_data_home"); - let cluster = GreptimeDbClusterBuilder::new("test_distributed_admin_function_events") - .await - .with_datanodes(1) - .with_store_config(store_config) - .with_shared_home_dir(Arc::new(home_dir)) - .build(true) - .await; - assert_admin_function_events(cluster.fe_instance()).await; -} - async fn assert_admin_function_events(instance: &Arc) { execute( instance, diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index c08b81211de..1c6d3f39633 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -21,6 +21,7 @@ use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -42,7 +43,7 @@ async fn test_event_table_auto_creation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -94,7 +95,7 @@ async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -185,6 +186,7 @@ async fn test_database_ddl_events() { let home_dir = create_temp_dir("test_database_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_frontend_auto_create_table(false) .with_store_config(store_config) @@ -208,6 +210,7 @@ async fn test_database_ddl_events() { async fn test_standalone_database_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_database_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 92888cced3c..41fe8382c82 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1453,9 +1453,14 @@ async fn metric_export_v2_cli_resume_roundtrip() { } #[tokio::test] -#[ignore = "requires GT_S3_* credentials and an S3-compatible test bucket"] async fn metric_export_v2_cli_s3_resume_roundtrip() { - metric_export_v2_cli_roundtrip(true).await; + // The CLI takes an explicit `--s3-endpoint`, so this only runs against an + // S3-compatible endpoint (MinIO in CI), not against bucket+region alone. + if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) + && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) + { + metric_export_v2_cli_roundtrip(true).await; + } } async fn metric_export_v2_cli_roundtrip(s3: bool) { diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 122c36c84f4..9d2884468d2 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -15,7 +15,9 @@ use std::sync::Arc; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use uuid::Uuid; use crate::event_recorder_test_util::{ @@ -31,6 +33,7 @@ const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; async fn test_standalone_flow_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_flow_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 4721a442a80..48749b784a0 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -43,6 +43,9 @@ async fn test_batch_gc_event() { common_telemetry::init_default_ut_logging(); let (store_config, _guard) = get_test_store_config(&store_type); let home_dir = create_temp_dir("test_batch_gc_event_data_home"); + // Keeps the production event-recorder flush interval: `assert_sst_count` + // counts SST files across every table, so eagerly flushed event rows would + // show up in the count. let cluster = GreptimeDbClusterBuilder::new("test_batch_gc_event") .await .with_datanodes(1) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 025d56b959a..970f65a2383 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -19,12 +19,13 @@ use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request as RequestBody; use api::v1::greptime_response::Response as ResponseBody; use api::v1::promql_request::Promql; +use api::v1::query_request::Query; use api::v1::value::ValueData; use api::v1::{ AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, ColumnSchema, CreateTableExpr, GreptimeRequest, InsertRequest, InsertRequests, PromInstantQuery, - PromRangeQuery, PromqlRequest, RequestHeader, Row, RowInsertRequest, RowInsertRequests, Rows, - SemanticType, Value, column, + PromRangeQuery, PromqlRequest, QueryRequest, RequestHeader, Row, RowInsertRequest, + RowInsertRequests, Rows, SemanticType, Value, column, }; use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -72,6 +73,7 @@ use tests_integration::test_util::{ setup_grpc_server_with_auto_create_table_disabled, setup_grpc_server_with_user_provider, }; use tonic::Request; +use tonic::codec::CompressionEncoding; use tonic::metadata::MetadataValue; use crate::both_deployment_cases; @@ -89,9 +91,13 @@ macro_rules! grpc_test { async fn [< $test >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { - let _ = $crate::grpc::$test(store_type).await; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::grpc::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); } - } )* } @@ -198,29 +204,128 @@ pub async fn test_grpc_message_size_ok(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Both the server (`servers::grpc::builder`) and the standard client +/// (`configure_tonic_client!`) enable zstd, so a plain round trip never shows +/// whether negotiation actually happened. This drives raw tonic clients instead +/// and asserts on the `grpc-encoding` the server answers with. pub async fn test_grpc_zstd_compression(store_type: StorageType) { - // server and client both support gzip - let config = GrpcServerConfig { - max_recv_message_size: 1024, - max_send_message_size: 1024, - ..Default::default() - }; - let (_db, fe_grpc_server) = setup_grpc_server_with( - store_type, - "test_grpc_zstd_compression", - None, - Some(config), - None, - ) - .await; + let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_grpc_zstd_compression").await; let addr = fe_grpc_server.bind_addr().unwrap().to_string(); - let grpc_client = Client::with_urls(vec![addr]); + let ddl = |sql: &str| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + })), + }; + + // Sends zstd and accepts zstd: the server has to decode a compressed request + // body and compress its response. + let mut zstd_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd) + .accept_compressed(CompressionEncoding::Zstd); + let response = zstd_client + .handle(Request::new(ddl( + "CREATE TABLE zstd_compression (ts TIMESTAMP TIME INDEX, payload STRING)", + ))) + .await + .unwrap(); + assert_eq!( + response + .metadata() + .get("grpc-encoding") + .map(|v| v.to_str().unwrap()), + Some("zstd") + ); + + // A payload far above the zstd frame overhead, so the request body really is + // compressed rather than passed through. + let payload = "compressible-".repeat(4096); + let insert = |ts: i64| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "zstd_compression".to_string(), + rows: Some(Rows { + schema: vec![ + ColumnSchema { + column_name: "ts".to_string(), + semantic_type: SemanticType::Timestamp as i32, + datatype: ColumnDataType::TimestampMillisecond as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "payload".to_string(), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::String as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(ts)), + }, + Value { + value_data: Some(ValueData::StringValue(payload.clone())), + }, + ], + }], + }), + }], + })), + }; + + let response = zstd_client + .handle(Request::new(insert(1000))) + .await + .unwrap(); + let ResponseBody::AffectedRows(rows) = response.into_inner().response.unwrap(); + assert_eq!(rows.value, 1); + + // A client that does not advertise zstd gets an uncompressed response. + let mut plain_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let response = plain_client + .handle(Request::new(insert(2000))) + .await + .unwrap(); + assert!(response.metadata().get("grpc-encoding").is_none()); + + // Both payloads survived their respective paths intact. let db = Database::new_with_dbname( format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), - grpc_client, + Client::with_urls(vec![addr]), ); - db.sql("show tables;").await.unwrap(); + let sql = format!( + "SELECT count(*) AS c FROM zstd_compression WHERE length(payload) = {}", + payload.len() + ); + let OutputData::Stream(stream) = db.sql(&sql).await.unwrap().data else { + panic!("expected a stream"); + }; + let recordbatches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + recordbatches.pretty_print().unwrap(), + "\ ++---+ +| c | ++---+ +| 2 | ++---+" + ); + let _ = fe_grpc_server.shutdown().await; } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 020e9c023c6..4cc409deec4 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -93,6 +93,31 @@ use yaml_rust::YamlLoader; use crate::both_deployment_cases; use crate::event_recorder_test_util::assert_procedure_actor_by_table; +/// Registers tests that exercise nothing below the HTTP layer, so running them +/// against a remote object store only pays bucket setup and teardown. +#[macro_export] +macro_rules! http_local_tests { + ($($(#[$meta:meta])* $test:ident),*,) => { + mod integration_http_local_test { + $( + #[tokio::test(flavor = "multi_thread")] + $( + #[$meta] + )* + async fn $test() { + let store_type = tests_integration::test_util::StorageType::File; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::http::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); + } + )* + } + }; +} + #[macro_export] macro_rules! http_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -134,12 +159,7 @@ macro_rules! http_tests { test_prometheus_promql_api, test_prometheus_label_replace_response, test_prom_http_api, - test_metrics_api, - test_health_api, - test_status_api, test_config_api, - test_dynamic_tracer_toggle, - test_dashboard_path, test_dashboard_api, test_prometheus_remote_write, test_prometheus_remote_write_v2, @@ -151,7 +171,6 @@ macro_rules! http_tests { test_vm_proto_remote_write, test_pipeline_api, - test_test_pipeline_api, test_pipeline_name_in_header, test_plain_text_ingestion, test_pipeline_auto_transform, @@ -187,8 +206,6 @@ macro_rules! http_tests { test_loki_json_logs_with_pipeline, test_elasticsearch_logs, test_elasticsearch_logs_with_index, - test_splunk_health, - test_splunk_health_is_public, test_splunk_logs, test_splunk_raw, test_log_query, @@ -1810,15 +1827,21 @@ pub async fn test_splunk_logs(store_type: StorageType) { .await; let client = TestClient::new(app).await; - // Authenticated SQL query (the user-provider harness requires auth on /v1/sql). + // Authenticated SQL statement (the user-provider harness requires auth on + // /v1/sql). Posted as a form so quotes and parentheses survive. async fn query(client: &TestClient, sql: &str) -> String { let res = client - .get(format!("/v1/sql?sql={sql}").as_str()) + .post("/v1/sql") .header("Authorization", basic_auth("greptime_user", "greptime_pwd")) + .form(&[("sql", sql)]) .send() .await; - assert_eq!(res.status(), StatusCode::OK, "query failed: {sql}"); - res.text().await + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "query failed: {sql}: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "query failed: {sql}: {body}"); + body } // HEC `Authorization: Splunk ` + JSON content type. @@ -1969,6 +1992,34 @@ transform: "custom pipeline should have dropped sourcetype (identity would keep it): {create}" ); + // 6b. Ingest into a pre-created table with a fulltext index, then flush: the + // index file is only written on flush, so this is the one path that reads it + // back. `matches_term` is the predicate the index applier recognises. + query( + &client, + "create table splunk_indexed (greptime_timestamp timestamp(9) time index, \ + \"event\" string fulltext index with(backend='bloom'))", + ) + .await; + let res = send_req( + &client, + splunk_headers(), + "/v1/splunk/services/collector/event?table=splunk_indexed", + br#"{"event":"disk pressure on web-01","time":1700000020}"#.to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + query(&client, "ADMIN FLUSH_TABLE('splunk_indexed')").await; + let rows = get_rows_from_output( + &query( + &client, + "select \"event\" from splunk_indexed where matches_term(\"event\", 'pressure')", + ) + .await, + ); + assert_eq!(rows, r#"[["disk pressure on web-01"]]"#); + // 7. Auth failures return HEC codes: missing token -> 2 (401), bad token -> 4 (403). let res = send_req( &client, @@ -2911,6 +2962,18 @@ pub async fn test_prometheus_remote_write(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "metric2").await; + validate_data( + "prometheus_remote_write_after_flush", + &client, + "select * from metric2", + table_val, + ) + .await; + // Prom RW tables carry the metric identity; type is inferred from naming. validate_data( "prometheus_remote_write_semantic_identity", @@ -3217,11 +3280,24 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + let histogram_sql = "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;"; + let histogram_rows = "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]"; validate_data( "prometheus_remote_write_v2_native_histogram_rows", &client, - "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + histogram_sql, + histogram_rows, + ) + .await; + + // The native histogram column is a struct of nested lists, so its SST + // encoding differs from its memtable representation. + flush_table(&client, "remote_write_v2_latency_seconds").await; + validate_data( + "prometheus_remote_write_v2_native_histogram_rows_after_flush", + &client, + histogram_sql, + histogram_rows, ) .await; @@ -4283,6 +4359,31 @@ transform: ) .await; + // The inverted, skipping and fulltext index files are only written on flush, + // so the queries below are the first thing that reads them back. + flush_table(&client, "test_db.logs1").await; + validate_data( + "pipeline_db_inverted_index", + &client, + "select id2 from test_db.logs1 where id1 = 2436", + "[[2528]]", + ) + .await; + validate_data( + "pipeline_db_skipping_index", + &client, + "select id1 from test_db.logs1 where type = 'I'", + "[[2436]]", + ) + .await; + validate_data( + "pipeline_db_fulltext_index", + &client, + "select id1 from test_db.logs1 where matches(log, 'ClusterAdapter')", + "[[2436]]", + ) + .await; + // 5. remove pipeline let encoded_ver_str: String = url::form_urlencoded::byte_serialize(version_str.as_bytes()).collect(); @@ -5079,6 +5180,24 @@ transform: ) .await; + // The fulltext and skipping index files are only written on flush, so the + // queries below are the first thing that reads them back. + flush_table(&client, "pipeline_index_options").await; + validate_data( + "pipeline_index_options_fulltext", + &client, + "select trace_id from pipeline_index_options where matches(message, 'greptime')", + "[[42]]", + ) + .await; + validate_data( + "pipeline_index_options_skipping", + &client, + "select message from pipeline_index_options where trace_id = 42", + r#"[["hello greptime"]]"#, + ) + .await; + guard.remove_all().await; } @@ -6995,6 +7114,18 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "claude_code_cost_usage_USD_total").await; + validate_data( + "otlp_metrics_all_select_after_flush", + &client, + "select * from `claude_code_cost_usage_USD_total` order by model desc;", + expected, + ) + .await; + // The synthesized resource descriptor: a plain mito info table keyed by // the raw OTel attribute keys (service.name is the only allowlisted // resource attribute in this payload), independent of the promote/ignore @@ -7846,6 +7977,17 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { ) .await; + // The attribute columns are JSON and the event/link columns are lists, whose + // SST encodings differ from their memtable representations. + flush_table(&client, "opentelemetry_traces").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from opentelemetry_traces;", + expected, + ) + .await; + guard.remove_all().await; } @@ -7972,6 +8114,30 @@ pub(crate) async fn test_otlp_traces_v2( ) .await; + // JSON2 path extraction has to resolve the same way against SST files as it + // does against the memtable. + flush_table(&client, "trace_v2_spans").await; + validate_data( + "otlp_traces_v2_json2_after_flush", + &client, + "select service_name, span_attributes.\"http.status_code\"::BIGINT, \ + resource_attributes.\"deployment.environment\"::STRING, \ + scope_attributes.enabled::BOOLEAN from trace_v2_spans;", + r#"[["frontend",200,"production",true]]"#, + ) + .await; + validate_data( + "otlp_traces_v2_events_links_after_flush", + &client, + r#"select json_get_string(span_events, '$[0].name'), + json_get_int(span_events, '$[0].attributes."event.code"'), + json_get_string(span_links, '$[0].trace_id'), + json_get_string(span_links, '$[0].attributes."link.type"') + from trace_v2_spans;"#, + r#"[["cache.hit",7,"cc9e0991a2e63d274984bd44ee669203","follows_from"]]"#, + ) + .await; + guard.remove_all().await; Ok(()) } @@ -8125,6 +8291,17 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { let expected = r#"[[1736480942444376000,1736480942444499000,123000,null,"c05d7a4ec8e1f231f02ed6e8da8655b4","d24f921c75f68e23","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444376000,1736480942444499000,123000,"d24f921c75f68e23","c05d7a4ec8e1f231f02ed6e8da8655b4","9630f2916e2f7909","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]],[1736480942444589000,1736480942444712000,123000,null,"cc9e0991a2e63d274984bd44ee669203","eba7be77e3558179","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444589000,1736480942444712000,123000,"eba7be77e3558179","cc9e0991a2e63d274984bd44ee669203","8f847259b0f6e1ab","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]]]"#; validate_data("otlp_traces", &client, "select * from mytable;", expected).await; + // `span_events` / `span_links` are JSON columns and the table is partitioned + // on `trace_id`, so the read back has to reassemble SST files across regions. + flush_table(&client, "mytable").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from mytable;", + expected, + ) + .await; + // The trace v1 main table carries the trace identity (events/links preserved as // JSON columns by the v1 model). validate_data( @@ -8961,6 +9138,17 @@ pub async fn test_otlp_logs(store_type: StorageType) { ) .await; + // The attribute columns are JSON, whose SST encoding differs from their + // memtable representation. + flush_table(&client, "opentelemetry_logs").await; + validate_data( + "otlp_logs_after_flush", + &client, + "select * from opentelemetry_logs;", + expected, + ) + .await; + // The auto-created log table carries the log identity. validate_data( "otlp_logs_semantic_identity", @@ -9343,6 +9531,17 @@ pub async fn test_loki_pb_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_pb_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9548,6 +9747,17 @@ pub async fn test_loki_json_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_json_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9767,7 +9977,7 @@ pub async fn test_log_query(store_type: StorageType) { // prepare data with SQL API let res = client - .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string);") + .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string fulltext index with(backend='bloom'));") .send() .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); @@ -9842,6 +10052,35 @@ pub async fn test_log_query(store_type: StorageType) { ] ); + // The fulltext index file is only written on flush. Replay the log query to + // cover the SST read path, and add a `matches_term` query, which is the + // predicate the index applier actually recognises. + flush_table(&client, "logs").await; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&log_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let resp = res.text().await; + let output = serde_json::from_str::(&resp).unwrap(); + let rows = output["output"][0]["records"]["rows"].as_array().unwrap(); + let mut flushed_messages = rows + .iter() + .map(|row| row[1].as_str().unwrap()) + .collect::>(); + flushed_messages.sort_unstable(); + assert_eq!(flushed_messages, messages); + + validate_data( + "log_query_fulltext_index", + &client, + "select message from logs where matches_term(message, 'at-explicit-end')", + r#"[["at-explicit-end"]]"#, + ) + .await; + let res = client .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);") .send() @@ -10865,11 +11104,20 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let mut req: ExportTraceServiceRequest = serde_json::from_str(content).unwrap(); - // Modify timestamp fields + // The fixture is written against a fixed base instant. Move it next to now: + // this table is created with `ttl=7d`, and stale spans would be dropped as + // expired the moment they are flushed into SST files. let now = Utc::now().timestamp_nanos_opt().unwrap() as u64; + let delta_micros = (now / 1_000) as i64 - 60_000_000 - JAEGER_V1_FIXTURE_BASE_MICROS as i64; + let delta_nanos = delta_micros * 1_000; for span in req.resource_spans.iter_mut() { for scope_span in span.scope_spans.iter_mut() { - // Only modify the timestamp fields for the span with the name "test-jaeger-get-operations" to current time. + for span in scope_span.spans.iter_mut() { + span.start_time_unix_nano = (span.start_time_unix_nano as i64 + delta_nanos) as u64; + span.end_time_unix_nano = (span.end_time_unix_nano as i64 + delta_nanos) as u64; + } + // `/api/operations` is queried without a time range, so these spans + // stay pinned to the last few seconds. if scope_span.scope.as_ref().unwrap().name == "test-jaeger-get-operations" { for span in scope_span.spans.iter_mut() { span.start_time_unix_nano = now - 5_000_000_000; // 5 seconds ago @@ -10946,7 +11194,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/operations` API. @@ -10995,12 +11244,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/services/{service_name}/operations` API. let res = client - .get("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11023,7 +11273,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API without start and end. @@ -11149,12 +11400,26 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); + assert_eq!(resp, expected); + + // Replay the same lookup against SST files: `trace_id` and `service_name` + // carry BLOOM skipping indexes whose puffin files only exist after a flush, + // and `span_events` / `span_links` are JSON columns. + flush_table(&client, trace_table_name).await; + let res = client + .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea") + .header("x-greptime-trace-table-name", trace_table_name) + .send() + .await; + assert_eq!(StatusCode::OK, res.status()); + let resp: Value = serde_json::from_str(&res.text().await).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API with start and end in microseconds. let res = client - .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11275,12 +11540,16 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API for non-existent trace. let res = client - .get("/v1/jaeger/api/traces/0000000000000000000000000000dead") + .get(&shift_fixture_micros( + "/v1/jaeger/api/traces/0000000000000000000000000000dead", + delta_micros, + )) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11299,12 +11568,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11426,13 +11696,14 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with tags. // 1. first query without tags, get 2 results let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11443,12 +11714,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // 2. second query with tags, get 1 result let res = client -.get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D") +.get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11459,14 +11731,15 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with Grafana User-Agent. // When user agent is Grafana, only return at most 3 spans per trace (earliest by timestamp). // Trace `5611dce1bc9ebed65352d99a027b08fb` has 4 spans, so only 3 should be returned. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Grafana/8.0.0") .send() @@ -11514,7 +11787,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API without User-Agent (default behavior). // All 4 spans should be returned for the trace. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11534,7 +11807,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API with Jaeger User-Agent (should return all spans like default). let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Jaeger-Query/1.0.0") .send() @@ -11613,6 +11886,59 @@ pub async fn test_influxdb_write(store_type: StorageType) { guard.remove_all().await; } +/// Flushes `table` so the rows written so far move from the memtable into SST +/// files. Replaying a query afterwards covers the encode/decode round trip and, +/// for indexed columns, the puffin files that only exist after a flush. +/// Base microsecond instant the `test_jaeger_query_api_for_trace_v1` fixture and +/// its expected responses are written against. +const JAEGER_V1_FIXTURE_BASE_MICROS: u64 = 1_738_726_754_492_421; + +/// Shifts every fixture timestamp in `text` by `delta` microseconds, so expected +/// responses and query bounds follow the payload after it is moved next to now. +/// +/// Only integer literals within a day of [`JAEGER_V1_FIXTURE_BASE_MICROS`] are +/// rewritten; durations, counts and ids are left alone. +fn shift_fixture_micros(text: &str, delta: i64) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(start) = rest.find(|c: char| c.is_ascii_digit()) { + out.push_str(&rest[..start]); + let digits_end = rest[start..] + .find(|c: char| !c.is_ascii_digit()) + .map_or(rest.len(), |offset| start + offset); + let digits = &rest[start..digits_end]; + match digits.parse::() { + Ok(value) if (value - JAEGER_V1_FIXTURE_BASE_MICROS as i64).abs() < 86_400_000_000 => { + out.push_str(&(value + delta).to_string()) + } + _ => out.push_str(digits), + } + rest = &rest[digits_end..]; + } + out.push_str(rest); + out +} + +async fn flush_table(client: &TestClient, table: &str) { + // The argument is parsed as a SQL table name, so an unquoted identifier is + // lower-cased and will not resolve a mixed-case table. + let quoted = table + .split('.') + .map(|part| format!("\"{part}\"")) + .collect::>() + .join("."); + let res = client + .post("/v1/sql") + .form(&[("sql", format!("ADMIN FLUSH_TABLE('{quoted}')"))]) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "flush {table} fail: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "flush {table} fail: {body}"); +} + async fn validate_data(test_name: &str, client: &TestClient, sql: &str, expected: &str) { let res = client .get(format!("/v1/sql?sql={sql}").as_str()) diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index b61b6cd3e19..278a7ba3851 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -53,6 +53,17 @@ grpc_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); http_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); +http_local_tests!( + test_health_api, + test_status_api, + test_metrics_api, + test_splunk_health, + test_splunk_health_is_public, + test_dynamic_tracer_toggle, + test_dashboard_path, + test_test_pipeline_api, +); + #[tokio::test(flavor = "multi_thread")] async fn test_http_auth_from_standalone_user_provider_config() { http::test_http_auth_from_standalone_user_provider_config().await; diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs index 19eaea6ff2c..a366217a7d9 100644 --- a/tests-integration/tests/reconciliation_event.rs +++ b/tests-integration/tests/reconciliation_event.rs @@ -20,6 +20,7 @@ use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use table::table_reference::TableReference; use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; @@ -33,6 +34,7 @@ async fn test_catalog_and_database_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("catalog_database_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -243,6 +245,7 @@ async fn test_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -351,6 +354,7 @@ async fn test_logical_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("logical_table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 906eb3cfee3..aaefcc2d49c 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -22,6 +22,7 @@ use mito2::gc::GcConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -44,6 +45,7 @@ async fn test_repartition_event() { let home_dir = create_temp_dir("test_repartition_event_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_repartition_event") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 8bad66dae3c..6cce1736151 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -41,7 +41,9 @@ use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use crate::event_recorder_test_util::{ assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, @@ -79,6 +81,7 @@ async fn test_table_ddl_procedure_events() { } let cluster = GreptimeDbClusterBuilder::new("table_ddl_procedure_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_metasrv_gc_config(gc_options) .with_datanode_gc_config(GcConfig { diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index b21838995c3..8d17e9b4e25 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -17,9 +17,9 @@ use std::sync::Arc; use common_test_util::temp_dir::create_temp_dir; use common_wal::config::DatanodeWalConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use uuid::Uuid; @@ -44,6 +44,7 @@ async fn test_view_ddl_events() { let home_dir = create_temp_dir("test_view_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_view_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) @@ -58,19 +59,6 @@ async fn test_view_ddl_events() { execute_view_ddl(instance, &source_table, &view).await; } -#[tokio::test(flavor = "multi_thread")] -async fn test_standalone_view_ddl_events() { - common_telemetry::init_default_ut_logging(); - let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_view_ddl_events") - .build() - .await; - let suffix = Uuid::new_v4().simple(); - let source_table = format!("view_ddl_event_source_{suffix}"); - let view = format!("view_ddl_event_{suffix}"); - - execute_view_ddl(standalone.fe_instance(), &source_table, &view).await; -} - async fn execute_view_ddl( instance: &Arc, source_table: &str, diff --git a/tests-integration/tests/wal_prune_event.rs b/tests-integration/tests/wal_prune_event.rs index db86ade4d9d..3879fb3e653 100644 --- a/tests-integration/tests/wal_prune_event.rs +++ b/tests-integration/tests/wal_prune_event.rs @@ -26,6 +26,7 @@ use meta_srv::procedure::wal_prune::{Context as WalPruneContext, WalPruneProcedu use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::record::Record; use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::assert_single_event; @@ -35,6 +36,7 @@ async fn test_standalone_wal_prune_event() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_wal_prune_event") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let topic_name = format!("test_standalone_wal_prune_event-{}", uuid::Uuid::new_v4()); From 0f118bc5ba3ad591824562f3ede4b8fd536fb86d Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 01:55:46 +0000 Subject: [PATCH 63/63] fix: serialize struct to json in postgres (#9170) * feat: serialize struct to json in postgres * fix: support view scalars and preserve null structs in scalar-to-value conversion Address PR review: - Utf8View/BinaryView ScalarValues now convert like their non-view forms instead of failing row extraction for struct columns - a null struct scalar converts to Value::Null so a null struct inside a list stays null in the serialized JSON Signed-off-by: Ning Sun * fix: return errors instead of panics for unsupported arrow field types Struct-typed query results with arrow field types greptimedb cannot represent (e.g. Decimal256) used to panic during schema conversion and row extraction, dropping the client connection. They now surface as query errors: - ConcreteDataType::try_from builds struct types fallibly via the new StructType::try_from_arrow_fields - Value::try_from(ScalarValue::Struct) uses the same fallible path - new try_value_from_array converts an arrow element to Value with error propagation, used by the postgres struct encoding Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/datatypes/src/data_type.rs | 4 +- src/datatypes/src/types/struct_type.rs | 19 ++ src/datatypes/src/value.rs | 99 +++++++- src/servers/src/postgres/types.rs | 315 ++++++++++++++++++++++++- tests-integration/src/test_util.rs | 55 +++++ tests-integration/tests/sql.rs | 233 +++++++++++++++++- 6 files changed, 711 insertions(+), 14 deletions(-) diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index f2ed05534d6..8ad388c2ff7 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -492,7 +492,9 @@ impl TryFrom<&ArrowDataType> for ConcreteDataType { ArrowDataType::Decimal128(precision, scale) => { ConcreteDataType::decimal128_datatype(*precision, *scale) } - ArrowDataType::Struct(fields) => ConcreteDataType::Struct(StructType::from(fields)), + ArrowDataType::Struct(fields) => { + ConcreteDataType::Struct(StructType::try_from_arrow_fields(fields)?) + } ArrowDataType::Float16 | ArrowDataType::Date64 | ArrowDataType::FixedSizeBinary(_) diff --git a/src/datatypes/src/types/struct_type.rs b/src/datatypes/src/types/struct_type.rs index ffd1d1c0b25..fdbfd4886b9 100644 --- a/src/datatypes/src/types/struct_type.rs +++ b/src/datatypes/src/types/struct_type.rs @@ -19,6 +19,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field}; use arrow_schema::Fields; use serde::{Deserialize, Serialize}; +use crate::error::Result; use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId}; use crate::value::Value; use crate::vectors::StructVectorBuilder; @@ -100,6 +101,24 @@ impl StructType { } } + /// Like the `From<&Fields>` impl, but returns an error for arrow field + /// types greptimedb cannot represent instead of panicking. + pub fn try_from_arrow_fields(fields: &Fields) -> Result { + let struct_fields = fields + .iter() + .map(|field| { + Ok(StructField::new( + field.name().clone(), + ConcreteDataType::try_from(field.data_type())?, + field.is_nullable(), + )) + }) + .collect::>>()?; + Ok(StructType { + fields: Arc::new(struct_fields), + }) + } + pub fn fields(&self) -> Arc> { self.fields.clone() } diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 56a551484f2..da1df882aba 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1152,12 +1152,13 @@ impl TryFrom for Value { ScalarValue::UInt16(u) => Value::from(u), ScalarValue::UInt32(u) => Value::from(u), ScalarValue::UInt64(u) => Value::from(u), - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => { + ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s) => { Value::from(s.map(StringBytes::from)) } ScalarValue::Binary(b) | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => Value::from(b.map(Bytes::from)), + | ScalarValue::FixedSizeBinary(_, b) + | ScalarValue::BinaryView(b) => Value::from(b.map(Bytes::from)), ScalarValue::List(array) => { // this is for item type let datatype = ConcreteDataType::try_from(&array.value_type())?; @@ -1221,7 +1222,15 @@ impl TryFrom for Value { .map(|v| Value::Decimal128(Decimal128::new(v, p, s))) .unwrap_or(Value::Null), ScalarValue::Struct(struct_array) => { - let struct_type = StructType::from(struct_array.fields()); + // A struct scalar carries a single element; a null (or empty) + // element means the struct itself is null, for example a null + // struct inside a list. + if struct_array.is_empty() || struct_array.is_null(0) { + return Ok(Value::Null); + } + // Build the struct type fallibly: an unrepresentable arrow + // field type must surface as an error, not a panic. + let struct_type = StructType::try_from_arrow_fields(struct_array.fields())?; let items = struct_array .columns() .iter() @@ -1244,8 +1253,6 @@ impl TryFrom for Value { | ScalarValue::LargeListView(_) | ScalarValue::Union(_, _, _) | ScalarValue::Float16(_) - | ScalarValue::Utf8View(_) - | ScalarValue::BinaryView(_) | ScalarValue::Map(_) | ScalarValue::Date64(_) | ScalarValue::RunEndEncoded(_, _, _) => { @@ -1259,6 +1266,21 @@ impl TryFrom for Value { } } +/// Converts the value at `index` of an arrow array into a [`Value`]. +/// +/// Unlike the unchecked conversion in [`crate::vectors::StructVector`]'s `get`, +/// this propagates conversion errors (for example arrow types greptimedb +/// cannot represent, such as `Decimal256`) instead of panicking, so callers +/// can surface a proper error to the user. +pub fn try_value_from_array(array: &dyn Array, index: usize) -> Result { + if array.is_null(index) { + return Ok(Value::Null); + } + let scalar = + ScalarValue::try_from_array(array, index).context(ConvertArrowArrayToScalarsSnafu)?; + Value::try_from(scalar) +} + impl From> for Value { fn from(value: ValueRef<'_>) -> Self { match value { @@ -2157,6 +2179,73 @@ pub(crate) mod tests { Value::Struct(struct_value), scalar_struct_value.try_into().unwrap() ); + + // view-typed strings and binaries convert like their non-view forms + assert_eq!( + Value::String("abc".into()), + ScalarValue::Utf8View(Some("abc".into())) + .try_into() + .unwrap() + ); + assert_eq!(Value::Null, ScalarValue::Utf8View(None).try_into().unwrap()); + assert_eq!( + Value::Binary(Bytes::from(vec![1, 2])), + ScalarValue::BinaryView(Some(vec![1, 2])) + .try_into() + .unwrap() + ); + assert_eq!( + Value::Null, + ScalarValue::BinaryView(None).try_into().unwrap() + ); + + // a null struct scalar is a null value, not a struct of null fields + let null_struct = ScalarStructBuilder::new_null(build_struct_type().as_arrow_fields()); + assert_eq!(Value::Null, null_struct.try_into().unwrap()); + } + + #[test] + fn test_try_value_from_array() { + use arrow_array::Int32Array; + use datafusion_common::arrow::datatypes::i256; + + let array = Int32Array::from(vec![Some(1), None]); + assert_eq!(Value::Int32(1), try_value_from_array(&array, 0).unwrap()); + assert_eq!(Value::Null, try_value_from_array(&array, 1).unwrap()); + + // a struct with representable fields converts + let supported = ScalarStructBuilder::new() + .with_name_and_scalar("i", ScalarValue::Int32(Some(7))) + .build() + .unwrap(); + let ScalarValue::Struct(array) = supported else { + unreachable!(); + }; + assert_eq!( + Value::Struct(StructValue::new( + vec![Value::Int32(7)], + StructType::new(std::sync::Arc::new(vec![StructField::new( + "i", + ConcreteDataType::int32_datatype(), + false, + )])) + )), + try_value_from_array(array.as_ref(), 0).unwrap() + ); + + // an unsupported arrow field type errors instead of panicking + let unsupported = ScalarStructBuilder::new() + .with_name_and_scalar( + "d", + ScalarValue::Decimal256(Some(i256::from_i128(1)), 38, 10), + ) + .build() + .unwrap(); + let ScalarValue::Struct(array) = unsupported else { + unreachable!(); + }; + let error = try_value_from_array(array.as_ref(), 0).unwrap_err(); + assert!(error.to_string().contains("Unsupported arrow data type")); } #[test] diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 033a1c665e2..c83debd39a1 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -19,7 +19,7 @@ use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; -use arrow::array::{Array, AsArray}; +use arrow::array::{Array, ArrayRef, AsArray}; use arrow_pg::encoder::{Encoder, encode_value}; use arrow_pg::list_encoder::encode_list; use arrow_schema::{DataType, TimeUnit}; @@ -37,7 +37,7 @@ use datatypes::json::JsonSettings; use datatypes::prelude::{ConcreteDataType, DataType as _, Value}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; -use datatypes::value::StructValue; +use datatypes::value::try_value_from_array; use futures::Stream; use pg_interval::Interval as PgInterval; use pgwire::api::Type; @@ -234,7 +234,8 @@ impl ToSqlText for OidAliasValue<'_> { } } -/// this function will encode greptime's `StructValue` into PostgreSQL jsonb type +/// this function will encode greptime's structured values (`StructValue`, or a +/// list of them) into PostgreSQL jsonb type /// /// Note that greptimedb has different types of StructValue for storing json data, /// based on policy defined in `JsonSettings`. But here the `StructValue` @@ -245,18 +246,35 @@ impl ToSqlText for OidAliasValue<'_> { /// arrays: element in array must be the same type fn encode_struct( _query_ctx: &QueryContextRef, - struct_value: StructValue, + value: Value, builder: &mut S, pg_field: &FieldInfo, ) -> PgWireResult<()> { let encoding_setting = JsonSettings::default(); let json_value = encoding_setting - .decode(Value::Struct(struct_value)) + .decode(value) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; builder.encode_field(&json_value, pg_field) } +/// Encodes one row of a structured column (`Struct`, or `List` of `Struct`) as +/// a PostgreSQL `json` value. +/// +/// Returns an error for arrow field types greptimedb cannot represent +/// (e.g. `Decimal256`) instead of panicking, so the client receives a proper +/// error response. +fn encode_structured_json( + query_ctx: &QueryContextRef, + column: &ArrayRef, + i: usize, + encoder: &mut S, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + let value = try_value_from_array(column.as_ref(), i).map_err(convert_err)?; + encode_struct(query_ctx, value, encoder, pg_field) +} + pub(crate) struct RecordBatchRowStream where S: Encoder, @@ -369,6 +387,11 @@ where } } + DataType::List(list_field) + if matches!(list_field.data_type(), DataType::Struct(_)) => + { + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; + } DataType::List(_) => { let array = column.as_list::(); let items = array.value(i); @@ -376,7 +399,7 @@ where encode_list(encoder, items, pg_field)?; } DataType::Struct(_) => { - encode_struct(query_ctx, Default::default(), encoder, pg_field)?; + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; } DataType::Utf8 => { let arrow_field = arrow_schema.field(j); @@ -445,7 +468,9 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC_ARRAY), &ConcreteDataType::Json(_) => Ok(Type::JSON_ARRAY), &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL_ARRAY), - &ConcreteDataType::Struct(_) => Ok(Type::JSON_ARRAY), + // a list of structs is encoded as a single JSON array, matching the + // struct-to-json mapping above + &ConcreteDataType::Struct(_) => Ok(Type::JSON), &ConcreteDataType::Dictionary(_) | &ConcreteDataType::Vector(_) | &ConcreteDataType::List(_) => server_error::UnsupportedDataTypeSnafu { @@ -1921,6 +1946,282 @@ mod test { } } + #[test] + fn test_encode_list_of_struct_as_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::{ListValue, StructValue}; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + // struct type: {a: int32, b: float64} + let fields = vec![ + StructField::new("a".to_string(), ConcreteDataType::int32_datatype(), true), + StructField::new("b".to_string(), ConcreteDataType::float64_datatype(), true), + ]; + let struct_type = + ConcreteDataType::struct_datatype(StructType::new(Arc::new(fields.clone()))); + let list_type = ConcreteDataType::list_datatype(Arc::new(struct_type.clone())); + + let make_struct = |a: i32, b: f64| { + Value::Struct(StructValue::new( + vec![Value::Int32(a), Value::Float64(b.into())], + StructType::new(Arc::new(fields.clone())), + )) + }; + let rows = vec![ + Value::List(ListValue::new( + vec![make_struct(1, 1.5), Value::Null, make_struct(3, 3.5)], + Arc::new(struct_type.clone()), + )), + Value::List(ListValue::new(vec![], Arc::new(struct_type.clone()))), + Value::Null, + ]; + let vector = Helper::try_from_row_into_vector(&rows, &list_type).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "l".to_string(), + list_type.clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + // The list-of-struct column is declared as a single json value. + let pg_schema = Arc::new(vec![FieldInfo::new( + "l".into(), + None, + None, + type_gt_to_pg(&list_type).unwrap(), + FieldFormat::Text, + )]); + assert_eq!(pg_schema[0].datatype(), &Type::JSON); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL list row stays NULL. + assert_eq!(decoded[2], None); + + // A list of structs is a JSON array; null elements stay null. + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!([ + {"a": 1, "b": 1.5}, + null, + {"a": 3, "b": 3.5}, + ]) + ); + + // An empty list is an empty JSON array. + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!(json, serde_json::json!([])); + } + + #[test] + fn test_encode_struct_value_as_json() { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use api::helper::pb_value_to_value_ref; + use api::v1::Value as PbValue; + use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, + }; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + fn histogram_value(histogram: &Histogram) -> Value { + let value_data = encode_native_histogram(histogram).unwrap(); + let pb_value = PbValue { + value_data: Some(value_data), + }; + let column_schema = native_histogram_column_schema().unwrap(); + Value::from(pb_value_to_value_ref( + &pb_value, + column_schema.datatype_extension.as_ref(), + )) + } + + // One integer-count histogram and one float-count histogram, mirroring + // what prometheus remote-write v2 stores. The expected JSON matches the + // HTTP query output for the same payloads. + let int_histogram = Histogram { + schema: 1, + sum: 10.0, + zero_threshold: 0.001, + reset_hint: 2, + start_timestamp: 1500, + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + count: Some(Count::CountInt(8)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }; + let float_histogram = Histogram { + schema: 2, + sum: 20.0, + zero_threshold: 0.002, + reset_hint: 3, + start_timestamp: 2500, + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + ..Default::default() + }; + + let rows = vec![ + histogram_value(&int_histogram), + histogram_value(&float_histogram), + Value::Null, + ]; + let vector = + Helper::try_from_row_into_vector(&rows, native_histogram_value_type()).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "greptime_native_histogram", + native_histogram_value_type().clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + let pg_schema = Arc::new(vec![FieldInfo::new( + "greptime_native_histogram".into(), + None, + None, + Type::JSON, + FieldFormat::Text, + )]); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL struct row stays NULL. + assert_eq!(decoded[2], None); + + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + } + #[test] fn test_encode_utf8_oid_alias_data() { let aliases = [ diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index fd02db2e5a1..37a1740ccd1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1264,6 +1264,61 @@ pub async fn setup_pg_server_with_user_provider( (instance.guard, Arc::new(pg_server)) } +/// Sets up a standalone instance with both a Prometheus remote-write HTTP app +/// (native histograms enabled) and a Postgres server attached, so native +/// histogram data written via remote-write can be queried over the Postgres +/// protocol. +pub async fn setup_pg_server_with_prom_native_histogram( + store_type: StorageType, + name: &str, +) -> (TestGuard, Router, Arc>) { + unsafe { + std::env::set_var("TZ", "UTC"); + } + + let instance = setup_standalone_instance(name, store_type).await; + + // Prometheus remote-write HTTP app with native histograms enabled. + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, true) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); + let app = http_server.build(http_server.make_app()).unwrap(); + + // Postgres server on the same instance. + let runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("pg-runtime") + .build() + .unwrap(); + + let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port()); + let opts = PostgresOptions { + addr: fe_pg_addr.clone(), + ..Default::default() + }; + let tls_server_config = Arc::new( + ReloadableTlsServerConfig::try_new(opts.tls.clone()) + .expect("Failed to load certificates and keys"), + ); + + let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + )); + + pg_server + .start(fe_pg_addr.parse::().unwrap()) + .await + .unwrap(); + + (instance.guard, app, Arc::new(pg_server)) +} + pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) { let catalog_manager = instance .catalog_manager() diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 6c58ab0246c..e3b66fc4447 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -31,7 +31,8 @@ use sqlx::{Connection, Executor, Row}; use tests_integration::test_util::{ StorageType, setup_mysql_server, setup_mysql_server_with_slow_query_threshold, setup_mysql_server_with_user_provider, setup_pg_server, - setup_pg_server_with_slow_query_threshold, setup_pg_server_with_user_provider, + setup_pg_server_with_prom_native_histogram, setup_pg_server_with_slow_query_threshold, + setup_pg_server_with_user_provider, }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; @@ -94,6 +95,8 @@ macro_rules! sql_tests { test_mysql_federated_prepare_stmt, test_mysql_prepare_tql_and_show, test_postgres_extended_query_row_returning_statements, + test_postgres_native_histogram, + test_postgres_struct_types, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -879,6 +882,234 @@ pub async fn test_postgres_bytea(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_postgres_native_histogram(store_type: StorageType) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use axum::http::StatusCode; + use prost::Message; + use servers::http::test_helpers::TestClient; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_store; + + let (mut guard, app, fe_pg_server) = + setup_pg_server_with_prom_native_histogram(store_type, "test_postgres_native_histogram") + .await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + let http_client = TestClient::new(app).await; + + // Ingest one integer-count and one float-count native histogram via + // prometheus remote-write v2. + let write_request = remote_write_v2::request_with_labels_and_histograms( + vec![ + (prom_store::METRIC_NAME_LABEL, "pg_native_histogram_seconds"), + ("job", "api"), + ("instance", "localhost:9090"), + ], + vec![ + Histogram { + count: Some(Count::CountInt(8)), + sum: 10.0, + schema: 1, + zero_threshold: 0.001, + zero_count: Some(ZeroCount::ZeroCountInt(1)), + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + reset_hint: 2, + timestamp: 3000, + start_timestamp: 1500, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + sum: 20.0, + schema: 2, + zero_threshold: 0.002, + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + reset_hint: 3, + timestamp: 4000, + start_timestamp: 2500, + ..Default::default() + }, + ], + ); + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = http_client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + let rows = client + .simple_query( + "select greptime_native_histogram from pg_native_histogram_seconds order by greptime_timestamp", + ) + .await + .unwrap(); + let jsons: Vec<&str> = rows + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + assert_eq!(jsons.len(), 2); + + // The struct column is encoded as a JSON object, matching the HTTP query + // output for the same histograms. Before the fix every row was encoded as + // an empty struct (`{}`). + let histogram: serde_json::Value = serde_json::from_str(jsons[0]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let histogram: serde_json::Value = serde_json::from_str(jsons[1]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + +pub async fn test_postgres_struct_types(store_type: StorageType) { + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_struct_types").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + async fn query_one(client: &Client, sql: &str) -> String { + let messages = client.simple_query(sql).await.unwrap(); + messages + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0).map(|v| v.to_string()), + _ => None, + }) + .next() + .unwrap() + } + + // Struct and list-of-struct columns are serialized as JSON. View-typed + // struct fields must not fail row extraction, and a null struct inside a + // list stays null instead of becoming a struct of null fields. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + // '0102' casts to its UTF-8 bytes when interpreted as BinaryView. + let row = query_one(&client, "SELECT struct(arrow_cast('0102', 'BinaryView'))").await; + assert_eq!(row, "{\"c0\":[48,49,48,50]}"); + + let row = query_one(&client, "SELECT struct([struct(1), NULL])").await; + assert_eq!(row, "{\"c0\":[{\"c0\":1},null]}"); + + // Top-level control: the same null element in a bare list stays null too. + let row = query_one(&client, "SELECT [struct(1), NULL]").await; + assert_eq!(row, "[{\"c0\":1},null]"); + + // Unsupported arrow field types surface as query errors instead of + // dropping the connection. + let error = client + .simple_query("SELECT struct(arrow_cast('1', 'Decimal256(38, 10)'))") + .await + .unwrap_err(); + let message = match error.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => error.to_string(), + }; + assert!( + message.contains("Unsupported arrow data type"), + "unexpected error message: {message}" + ); + + // The connection stays usable after the error. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_slow_query(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server_with_slow_query_threshold( store_type,