From 25bca4609f99fdd46078c528e4ad84aea7f7281f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 11:25:50 +0000 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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 05/10] 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 06/10] 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 07/10] 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 08/10] 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 09/10] 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 10/10] 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,