mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-21 11:22:31 +00:00
fix: backport suggested independent fixes to release/5e02efa9 (#9795)
* fix(servers): resolve port from IPv6 bind addresses (#9142) Signed-off-by: immanuwell <pchpr.00@list.ru> (cherry picked from commit5ec679fed9) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: register region migration procedures before returning IDs (#9163) Signed-off-by: jeremyhi <fengjiachun@gmail.com> (cherry picked from commitc65e8d0962) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(promql): skip empty batches in SeriesDivide stream (#9178) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit5e57b00a95) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(servers): merge duplicate timeseries across RecordBatches in Prometheus remote read (#9032) * fix(servers): merge timeseries across RecordBatches in Prometheus remote read Signed-off-by: Tyagiquamar <mohdquamartyagi@gmail.com> * fix(servers): exercise recordbatches_to_timeseries and sort samples with sort_unstable_by_key --------- Signed-off-by: Tyagiquamar <mohdquamartyagi@gmail.com> Co-authored-by: dennis zhuang <killme2008@gmail.com> (cherry picked from commit59eb63c663) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(prometheus): honor label matchers in __name__ values query (#9134) * fix(prometheus): honor label matchers in __name__ values query `/api/v1/label/__name__/values?match[]={pod="abc"}` dropped every matcher other than `__name__` and returned all metrics in the schema. No error, just the wrong list. Grafana's metrics browser sends this request, so picking a label value there did nothing. Selectors that only constrain `__name__` keep answering from table metadata. A selector constraining an ordinary label now goes to the data: scan each metric engine physical table for distinct `__table_id` in the time range, map the ids back to metric names, then apply the selector's own `__name__` matchers. Only metric engine tables are covered; other engines share no column space to scan. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * refactor(prometheus): batch-resolve metric names by table id Building a full table-id-to-name map meant walking every table in the schema and holding all of them in memory, just to name the handful the scan returned. Use `tables_by_ids` instead — one batch KV read over the ids the scan actually produced. The catalog walk stays, but only to find the physical tables to scan. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(promql): read an absent label as the empty string A matcher on a label the series does not carry only worked when the table had no column for it at all. Where the column exists but is NULL on that row -- the norm for logical metrics sharing a metric engine physical table, which holds the union of their label columns -- three-valued logic dropped the row, so `host!="host1"` and `host=""` missed every metric without a host label. Coalesce nullable string label columns to "" for matchers that accept the empty string, rather than only for the OTLP temporality marker. Equality matchers are untouched; they cannot match NULL either way. This is the Prometheus compatibility fix #8970 deliberately kept out of its own scope. The cost is visible in the regex sqlness plan: the predicate becomes a CASE, so the scan loses its LastRow selector and grows a FilterExec. Only negative and empty-accepting matchers pay it. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(promql): don't panic on pre-epoch label value bounds `rewrite_label_values_query` unwrapped `duration_since(UNIX_EPOCH)`, which returns an error for an instant before the epoch. `start=1969-12-31T23:59:59Z` parses as valid RFC3339, so the request panicked instead of answering. Recover the sign from the error branch, and report a value beyond i64 milliseconds as an error rather than wrapping the cast. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * refactor(prometheus): drop applicable_matchers, share the distinct scan With the planner reading an absent label as empty, the frontend no longer needs to pre-filter matchers per physical table. Removing that exposed a second problem: a physical table that never took a column from a logical table exposes no `__table_id`, and projecting it failed the whole request. Skip those tables; the only thing that can miss is a metric with no labels. Also pulls out the plan-build-execute-collect sequence the two label value scans had in common. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit007836ba7a) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(mito2): cancel cache construction for incomplete scans (#9254) * fix(mito2): cancel cache construction for incomplete scans CacheBatchBuffer spawned the background concat task and dropped its join handle. A scan that was cancelled or failed therefore left the task alive: it kept compacting already queued batches, and while waiting for a range result memory permit it held them, even though without a finish command the result can never be put into the cache. Keep the handle and abort it when the buffer is dropped. The handle is cleared once the task owns the finish command, so a completed scan still populates the cache after its stream is dropped. Abort does not preempt a concat that is already running; it takes effect the next time the task is polled. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(mito2): wait for the permit park before cancelling the buffer An empty buffered_batches only proves the batches were enqueued, so the cancellation test could abort a concat task that had never been polled. Count acquisitions that find too few permits, a test-only signal, and drop the buffer once the task has reached that wait. Nothing awaits between the check and the parking, so an observed increment means the caller is about to wait for permits the test holds for the rest of the case. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit43ce07fa69) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: merge every state row in geo_path and json_encode_path (#9339) Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commitcc1be82858) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: align ordered aggregate state type with the accumulator output (#9340) * fix: align ordered aggregate state type with the accumulator output Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix: reject mismatched aggregate states and keep hard-ordered aggregates unsplit Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix: keep WITHIN GROUP aggregates splittable and relabel state fields by position Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commitaffc0a1b1d) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(mito2): stop parallel flat scan tasks once the receiver is dropped (#9368) * fix(mito2): stop parallel flat scan tasks once the receiver is dropped spawn_flat_scan_task ignored send errors, so after a query was cancelled the detached task kept reading its source to EOF. Break out of the loop when the receiver is gone. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * chore(mito2): log when a parallel scan task stops on a dropped receiver Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * chore(mito2): include the send error in the parallel scan stop log Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit411c0f173e) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(catalog): skip backend batch_get on full cache hits (#9431) fix(catalog): return cached batch values without an empty backend call Count backend batch_get invocations in regression tests and cover empty input. Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commitacebd6b35b) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: coalesce equal histogram quantile bucket boundaries (#9451) * fix: coalesce equal histogram quantile bucket boundaries Merge numerically equivalent bucket boundaries before monotonicity repair, preserving the allocation-free path for unique boundaries. Share coalescing with histogram_fraction and cover duplicate boundaries with unit and SQL regressions. Fixes #9443 Signed-off-by: evenyag <realevenyag@gmail.com> * fix: keep duplicate infinity buckets in one histogram group Confirm classic histogram completion at a series/timestamp boundary or EOF instead of the first positive infinity. Retain lookahead across batches and reject incomplete optimistic groups so equivalent infinity labels reach numeric coalescing together. Add stream and SQL regressions for duplicate infinity labels, including established bucket counts and EOF. Correct older fixtures to use distinct keys for distinct histogram groups. Signed-off-by: evenyag <realevenyag@gmail.com> * fix: avoid buffering histogram groups with invalid trailing bounds Discover candidate bucket counts at the first group boundary and let optimistic validation fall back to safe mode for malformed bounds. Document discovery and cover trailing invalid and null labels with a focused streaming regression. Signed-off-by: evenyag <realevenyag@gmail.com> --------- Signed-off-by: evenyag <realevenyag@gmail.com> (cherry picked from commit97913b4b3a) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(otlp): coerce empty and "-" trace attribute strings to null (#9461) * fix(otlp): coerce empty and "-" trace attribute strings to null Instrumentations fill numeric attributes from absent headers with an empty string (PSR-7 getHeaderLine returns ""). Coercing that into an existing numeric column failed and rejected the whole span. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test: trim redundant row from empty-string rewrite test Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commitd640ab8bc5) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(metasrv): retry failed heartbeat address updates (#9437) * fix(metasrv): retry failed heartbeat address updates Signed-off-by: WenyXu <wenymedia@gmail.com> * refactor(metasrv): move shared heartbeat helpers to utils Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com> (cherry picked from commit65b9a35692) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: align TWCS assignment with half-open time windows (#9464) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> (cherry picked from commit5126623e63) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(promql): keep vector labels when a computed scalar is on the right (#9478) * fix(promql): keep vector labels when a computed scalar is on the right `metric + time()` or `metric * scalar(other)` planned the join with the scalar side's context, which has no tags, so every series collapsed into one unlabeled series. Use the vector side's context when the right operand is scalar-typed. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(promql): update results for labels kept with a computed scalar on the right v * pi() and v + time() now keep the vector's labels and time index name, as in Prometheus. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(promql): qualify columns of an unnamed operand in vector joins A comparison filter's projection has no table name, so after the join its time index resolved unqualified and was ambiguous when the other operand had the same column. Also keep the empty-metric branch unchanged for scalar left operands. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(promql): cover comparison filters joined with a scalar Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit6da5352c53) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(promql): keep metric name for comparison operators (#9332) * fix(promql): keep metric name for comparison operators `collect_metric_names` dropped the response `__name__` for every binary operator except `and` / `or` / `unless`, so comparisons such as `node_load1 > 5` came back without the metric name. Grafana rendered them as `{instance="..."}` and alert templates reading `$labels.__name__` got nothing. Follow Prometheus' `shouldDropMetricName` / `resultMetric`: arithmetic operators drop the name, comparisons keep the left-hand name, except when the comparison returns a bool or is a one-to-one `on(...)` match. A `group_left` / `group_right` modifier does not reduce the labels, so the name is kept there. Both the HTTP PromQL API and the gRPC gateway share `retrieve_metric_name_and_result_type`, so one fix covers both. Signed-off-by: Aumnertic <hznuzc@gmail.com> * docs: document binary_keeps_metric_name Signed-off-by: Aumnertic <hznuzc@gmail.com> * fix(promql): handle on(__name__), scalar lhs and group_right Signed-off-by: Aumnertic <hznuzc@gmail.com> * chore(promql): correct result-side docs and drop unreachable String arm Signed-off-by: Aumnertic <hznuzc@gmail.com> * fix(promql): drop metric name when explicitly ignored Handle ignoring(__name__) for one-to-one comparisons and cover comparison, group_left, group_right and unless behavior. Signed-off-by: Aumnertic <hznuzc@gmail.com> --------- Signed-off-by: Aumnertic <hznuzc@gmail.com> (cherry picked from commitc33a6cfa56) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(query): keep filter fetch when rewriting matches term (#9521) * test: cover matches term limit dropped by constant rewrite Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep filter fetch when rewriting matches term Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit483d1e49a3) Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * style: format #9437 test imports for the release toolchain Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * test: regenerate array_agg result for release plan rendering Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> --------- Signed-off-by: immanuwell <pchpr.00@list.ru> Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> Signed-off-by: jeremyhi <fengjiachun@gmail.com> Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: Tyagiquamar <mohdquamartyagi@gmail.com> Signed-off-by: Dennis Zhuang <killme2008@gmail.com> Signed-off-by: WenyXu <wenymedia@gmail.com> Signed-off-by: evenyag <realevenyag@gmail.com> Signed-off-by: Aumnertic <hznuzc@gmail.com> Co-authored-by: Immanuel Tikhonov <122638311+immanuwell@users.noreply.github.com> Co-authored-by: jeremyhi <jiachun_feng@proton.me> Co-authored-by: discord9 <discord9@163.com> Co-authored-by: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Co-authored-by: dennis zhuang <killme2008@gmail.com> Co-authored-by: Weny Xu <wenymedia@gmail.com> Co-authored-by: Yingwen <realevenyag@gmail.com> Co-authored-by: Aum <hznuzc@gmail.com>
This commit is contained in:
Lei, HUANG
co-authored by
Immanuel Tikhonov
jeremyhi
discord9
Mohd Quamar Tyagi
dennis zhuang
Weny Xu
Yingwen
Aum
Immanuel Tikhonov
jeremyhi
discord9
Mohd Quamar Tyagi
dennis zhuang
Weny Xu
Yingwen
Aum
parent
fd3faeb7ed
commit
2b05da4c06
@@ -192,6 +192,10 @@ impl KvBackend for CachedKvBackend {
|
||||
}
|
||||
}
|
||||
|
||||
if miss_keys.is_empty() {
|
||||
return Ok(BatchGetResponse { kvs });
|
||||
}
|
||||
|
||||
let batch_get_req = BatchGetRequest::new().with_keys(miss_keys.clone());
|
||||
|
||||
let pre_version = self.version();
|
||||
@@ -503,6 +507,7 @@ mod tests {
|
||||
pub struct SimpleKvBackend {
|
||||
inner_map: DashMap<Vec<u8>, Vec<u8>>,
|
||||
get_execute_times: Arc<AtomicU32>,
|
||||
batch_get_execute_times: Arc<AtomicU32>,
|
||||
}
|
||||
|
||||
impl TxnService for SimpleKvBackend {
|
||||
@@ -520,6 +525,7 @@ mod tests {
|
||||
}
|
||||
|
||||
async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse, Self::Error> {
|
||||
self.batch_get_execute_times.fetch_add(1, Ordering::SeqCst);
|
||||
let mut kvs = Vec::with_capacity(req.keys.len());
|
||||
for key in req.keys.iter() {
|
||||
if let Some(kv) = self.get(key).await? {
|
||||
@@ -571,6 +577,7 @@ mod tests {
|
||||
async fn test_cached_kv_backend() {
|
||||
let simple_kv = Arc::new(SimpleKvBackend::default());
|
||||
let get_execute_times = simple_kv.get_execute_times.clone();
|
||||
let batch_get_execute_times = simple_kv.batch_get_execute_times.clone();
|
||||
let cached_kv = CachedKvBackend::wrap(simple_kv);
|
||||
|
||||
add_some_vals(&cached_kv).await;
|
||||
@@ -582,9 +589,11 @@ mod tests {
|
||||
assert_eq!(get_execute_times.load(Ordering::SeqCst), 0);
|
||||
|
||||
for _ in 0..10 {
|
||||
let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap();
|
||||
let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap();
|
||||
|
||||
assert_eq!(batch_get_resp.kvs.len(), 2);
|
||||
assert_eq!(get_execute_times.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
let batch_get_req = BatchGetRequest {
|
||||
@@ -594,14 +603,29 @@ mod tests {
|
||||
let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap();
|
||||
|
||||
assert_eq!(get_execute_times.load(Ordering::SeqCst), 3);
|
||||
assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2);
|
||||
|
||||
for _ in 0..10 {
|
||||
let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap();
|
||||
let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap();
|
||||
|
||||
assert_eq!(batch_get_resp.kvs.len(), 3);
|
||||
assert_eq!(get_execute_times.load(Ordering::SeqCst), 3);
|
||||
assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cached_kv_backend_empty_batch_get() {
|
||||
let simple_kv = Arc::new(SimpleKvBackend::default());
|
||||
let cached_kv = CachedKvBackend::wrap(simple_kv.clone());
|
||||
|
||||
let response = cached_kv.batch_get(BatchGetRequest::new()).await.unwrap();
|
||||
|
||||
assert!(response.kvs.is_empty());
|
||||
assert_eq!(simple_kv.batch_get_execute_times.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(simple_kv.get_execute_times.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cached_kv_backend_rejects_writes_with_read_only_inner() {
|
||||
let inner = Arc::new(MemoryKvBackend::<common_meta::error::Error>::new());
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::array::{ArrayRef, BooleanArray, StructArray};
|
||||
use arrow::array::{ArrayData, ArrayRef, BooleanArray, StructArray, make_array};
|
||||
use arrow_schema::{FieldRef, Fields};
|
||||
use common_telemetry::debug;
|
||||
use datafusion::functions_aggregate::all_default_aggregate_functions;
|
||||
@@ -84,6 +84,19 @@ pub fn is_all_aggr_exprs_steppable(aggr_exprs: &[Expr]) -> bool {
|
||||
return false;
|
||||
}
|
||||
|
||||
// DataFusion only sorts the input of an aggregate with a hard ordering requirement
|
||||
// when the requirement is already satisfied or the aggregate has a reverse
|
||||
// expression (apache/datafusion#25676). The state wrapper has none, so e.g.
|
||||
// `nth_value(.. ORDER BY ..)` would read unsorted input on datanodes. Ordered-set
|
||||
// aggregates like `approx_percentile_cont(..) WITHIN GROUP (ORDER BY ..)` are
|
||||
// exempt: their ORDER BY names the value, and they don't need sorted input.
|
||||
if !aggr_func.params.order_by.is_empty()
|
||||
&& aggr_func.func.order_sensitivity().hard_requires()
|
||||
&& !aggr_func.func.supports_within_group_clause()
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// whether the corresponding state function exists in the registry
|
||||
FUNCTION_REGISTRY.is_aggr_func_exist(&aggr_state_func_name(aggr_func.func.name()))
|
||||
} else {
|
||||
@@ -516,43 +529,7 @@ impl StateGroupsAccum {
|
||||
}
|
||||
|
||||
fn wrap_state_arrays(&self, arrays: Vec<ArrayRef>) -> datafusion_common::Result<ArrayRef> {
|
||||
let array_type = arrays
|
||||
.iter()
|
||||
.map(|array| array.data_type().clone())
|
||||
.collect::<Vec<_>>();
|
||||
let expected_type = self
|
||||
.state_fields
|
||||
.iter()
|
||||
.map(|field| field.data_type().clone())
|
||||
.collect::<Vec<_>>();
|
||||
if array_type != expected_type {
|
||||
debug!(
|
||||
"State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}",
|
||||
self.state_fields.len(),
|
||||
arrays.len(),
|
||||
self.state_fields,
|
||||
array_type,
|
||||
);
|
||||
let guess_schema = arrays
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, array)| {
|
||||
Field::new(
|
||||
format!("col_{index}[mismatch_state]").as_str(),
|
||||
array.data_type().clone(),
|
||||
true,
|
||||
)
|
||||
})
|
||||
.collect::<Fields>();
|
||||
let array = StructArray::try_new(guess_schema, arrays, None)?;
|
||||
return Ok(Arc::new(array));
|
||||
}
|
||||
|
||||
Ok(Arc::new(StructArray::try_new(
|
||||
self.state_fields.clone(),
|
||||
arrays,
|
||||
None,
|
||||
)?))
|
||||
Ok(Arc::new(state_struct_array(&self.state_fields, arrays)?))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -605,6 +582,81 @@ impl GroupsAccumulator for StateGroupsAccum {
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps the state arrays of an accumulator into a struct of the declared state fields.
|
||||
///
|
||||
/// The declared state type is derived from logical expressions, while the accumulator names
|
||||
/// nested fields after physical expressions. For example, `array_agg(v ORDER BY ts)` declares
|
||||
/// its orderings as `List(Struct("ts": ..))` but produces `List(Struct("ts@0": ..))`. Arrays that
|
||||
/// differ only in nested field names are cast to the declared type; any other difference is an
|
||||
/// error.
|
||||
fn state_struct_array(
|
||||
state_fields: &Fields,
|
||||
arrays: Vec<ArrayRef>,
|
||||
) -> datafusion_common::Result<StructArray> {
|
||||
if arrays.len() != state_fields.len() {
|
||||
return Err(datafusion_common::DataFusionError::Internal(format!(
|
||||
"Expected {} state arrays for fields {:?}, got {}",
|
||||
state_fields.len(),
|
||||
state_fields,
|
||||
arrays.len()
|
||||
)));
|
||||
}
|
||||
let arrays = arrays
|
||||
.into_iter()
|
||||
.zip(state_fields.iter())
|
||||
.map(|(array, field)| {
|
||||
let expected = field.data_type();
|
||||
if array.data_type() == expected {
|
||||
Ok(array)
|
||||
} else if array.data_type().equals_datatype(expected) {
|
||||
Ok(make_array(relabel_nested_fields(
|
||||
array.to_data(),
|
||||
expected,
|
||||
)?))
|
||||
} else {
|
||||
Err(datafusion_common::DataFusionError::Internal(format!(
|
||||
"State field `{}` expects type {expected}, but the accumulator produced {}",
|
||||
field.name(),
|
||||
array.data_type()
|
||||
)))
|
||||
}
|
||||
})
|
||||
.collect::<datafusion_common::Result<Vec<_>>>()?;
|
||||
Ok(StructArray::try_new(state_fields.clone(), arrays, None)?)
|
||||
}
|
||||
|
||||
/// Rebuilds `data` with the type `target`, which must match it position by position apart
|
||||
/// from nested field names and metadata. Unlike a cast, children are never matched by name.
|
||||
fn relabel_nested_fields(
|
||||
data: ArrayData,
|
||||
target: &DataType,
|
||||
) -> datafusion_common::Result<ArrayData> {
|
||||
let child_types = match target {
|
||||
DataType::List(field)
|
||||
| DataType::LargeList(field)
|
||||
| DataType::FixedSizeList(field, _)
|
||||
| DataType::Map(field, _) => vec![field.data_type()],
|
||||
DataType::Struct(fields) => fields.iter().map(|f| f.data_type()).collect(),
|
||||
_ if data.child_data().is_empty() => vec![],
|
||||
_ => {
|
||||
return Err(datafusion_common::DataFusionError::NotImplemented(format!(
|
||||
"Relabeling nested fields of {target}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.zip(child_types)
|
||||
.map(|(child, child_type)| relabel_nested_fields(child.clone(), child_type))
|
||||
.collect::<datafusion_common::Result<Vec<_>>>()?;
|
||||
Ok(data
|
||||
.into_builder()
|
||||
.data_type(target.clone())
|
||||
.child_data(children)
|
||||
.build()?)
|
||||
}
|
||||
|
||||
impl StateAccum {
|
||||
pub fn new(
|
||||
inner: Box<dyn Accumulator>,
|
||||
@@ -631,40 +683,7 @@ impl Accumulator for StateAccum {
|
||||
.iter()
|
||||
.map(|s| s.to_array())
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let array_type = array
|
||||
.iter()
|
||||
.map(|a| a.data_type().clone())
|
||||
.collect::<Vec<_>>();
|
||||
let expected_type: Vec<_> = self
|
||||
.state_fields
|
||||
.iter()
|
||||
.map(|f| f.data_type().clone())
|
||||
.collect();
|
||||
if array_type != expected_type {
|
||||
debug!(
|
||||
"State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}",
|
||||
self.state_fields.len(),
|
||||
array.len(),
|
||||
self.state_fields,
|
||||
array_type,
|
||||
);
|
||||
let guess_schema = array
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, array)| {
|
||||
Field::new(
|
||||
format!("col_{index}[mismatch_state]").as_str(),
|
||||
array.data_type().clone(),
|
||||
true,
|
||||
)
|
||||
})
|
||||
.collect::<Fields>();
|
||||
let arr = StructArray::try_new(guess_schema, array, None)?;
|
||||
|
||||
return Ok(ScalarValue::Struct(Arc::new(arr)));
|
||||
}
|
||||
|
||||
let struct_array = StructArray::try_new(self.state_fields.clone(), array, None)?;
|
||||
let struct_array = state_struct_array(&self.state_fields, array)?;
|
||||
Ok(ScalarValue::Struct(Arc::new(struct_array)))
|
||||
}
|
||||
|
||||
|
||||
@@ -154,13 +154,17 @@ fn rewrite_expr(
|
||||
if is_fix {
|
||||
// then always fix the ordering field&distinct flag and more
|
||||
let order_by = aggregate_function.params.order_by.clone();
|
||||
// DataFusion always makes ordering fields nullable when building the physical
|
||||
// aggregate (`ordering_fields` in `datafusion-functions-aggregate-common`), and some
|
||||
// accumulators (e.g. `array_agg`) nest these fields in their state type. Keep the
|
||||
// declared state type in line with what the accumulator produces.
|
||||
let ordering_fields: Vec<_> = order_by
|
||||
.iter()
|
||||
.map(|sort_expr| {
|
||||
sort_expr
|
||||
.expr
|
||||
.to_field(&aggregate_input.schema())
|
||||
.map(|(_, f)| f)
|
||||
.map(|(_, f)| Arc::new(f.as_ref().clone().with_nullable(true)))
|
||||
})
|
||||
.collect::<datafusion_common::Result<Vec<_>>>()?;
|
||||
let distinct = aggregate_function.params.distinct;
|
||||
|
||||
@@ -27,8 +27,11 @@ use common_telemetry::init_default_ut_logging;
|
||||
use datafusion::catalog::{Session, TableProvider};
|
||||
use datafusion::datasource::DefaultTableSource;
|
||||
use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
|
||||
use datafusion::functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf;
|
||||
use datafusion::functions_aggregate::array_agg::array_agg_udaf;
|
||||
use datafusion::functions_aggregate::average::avg_udaf;
|
||||
use datafusion::functions_aggregate::count::count_udaf;
|
||||
use datafusion::functions_aggregate::nth_value::nth_value_udaf;
|
||||
use datafusion::functions_aggregate::sum::sum_udaf;
|
||||
use datafusion::optimizer::AnalyzerRule;
|
||||
use datafusion::optimizer::analyzer::type_coercion::TypeCoercion;
|
||||
@@ -1312,6 +1315,40 @@ async fn test_udaf_correct_eval_result() {
|
||||
order_by: vec![],
|
||||
null_treatment: None,
|
||||
},
|
||||
// The ordering state of `array_agg` nests the ORDER BY fields in `List(Struct(..))`.
|
||||
TestCase {
|
||||
func: array_agg_udaf(),
|
||||
input_schema: Arc::new(arrow_schema::Schema::new(vec![
|
||||
Field::new("number", DataType::Float64, true),
|
||||
Field::new(
|
||||
"ts",
|
||||
DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
])),
|
||||
args: vec![Expr::Column(Column::new_unqualified("number"))],
|
||||
input: vec![
|
||||
Arc::new(Float64Array::from(vec![Some(3.), Some(1.), Some(2.)])),
|
||||
Arc::new(TimestampMillisecondArray::from(vec![3000, 1000, 2000])),
|
||||
],
|
||||
expected_output: Some(ScalarValue::List(ScalarValue::new_list_nullable(
|
||||
&[
|
||||
ScalarValue::Float64(Some(1.)),
|
||||
ScalarValue::Float64(Some(2.)),
|
||||
ScalarValue::Float64(Some(3.)),
|
||||
],
|
||||
&DataType::Float64,
|
||||
))),
|
||||
expected_fn: None,
|
||||
distinct: false,
|
||||
filter: None,
|
||||
order_by: vec![SortExpr::new(
|
||||
Expr::Column(Column::new_unqualified("ts")),
|
||||
true,
|
||||
true,
|
||||
)],
|
||||
null_treatment: None,
|
||||
},
|
||||
// TODO(discord9): udd_merge/hll_merge/geo_path/quantile_aggr tests
|
||||
];
|
||||
let test_table_ref = TableReference::bare("TestTable");
|
||||
@@ -1421,3 +1458,89 @@ async fn execute_phy_plan(
|
||||
}
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_state_struct_array_rejects_mismatched_state() {
|
||||
let fields = Fields::from(vec![Field::new("sum", DataType::Int64, true)]);
|
||||
let arrays: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![1.0]))];
|
||||
let err = state_struct_array(&fields, arrays).unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("State field `sum` expects type Int64, but the accumulator produced Float64"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_state_struct_array_keeps_child_order() {
|
||||
let int_field = |name: &str| Field::new(name, DataType::Int64, true);
|
||||
// Same child types, names crossed: children must stay in place, not be matched by name.
|
||||
let produced = StructArray::from(vec![
|
||||
(
|
||||
Arc::new(int_field("a")),
|
||||
Arc::new(Int64Array::from(vec![10])) as ArrayRef,
|
||||
),
|
||||
(
|
||||
Arc::new(int_field("b")),
|
||||
Arc::new(Int64Array::from(vec![20])) as ArrayRef,
|
||||
),
|
||||
]);
|
||||
let declared = Fields::from(vec![Field::new(
|
||||
"state",
|
||||
DataType::Struct(Fields::from(vec![int_field("b"), int_field("a")])),
|
||||
true,
|
||||
)]);
|
||||
|
||||
let state = state_struct_array(&declared, vec![Arc::new(produced)]).unwrap();
|
||||
let state = state.column(0).as_struct();
|
||||
assert_eq!(
|
||||
state
|
||||
.column_by_name("b")
|
||||
.unwrap()
|
||||
.as_primitive::<arrow::datatypes::Int64Type>()
|
||||
.value(0),
|
||||
10
|
||||
);
|
||||
assert_eq!(
|
||||
state
|
||||
.column_by_name("a")
|
||||
.unwrap()
|
||||
.as_primitive::<arrow::datatypes::Int64Type>()
|
||||
.value(0),
|
||||
20
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hard_ordered_aggr_not_steppable() {
|
||||
let order_by = vec![SortExpr::new(
|
||||
Expr::Column(Column::new_unqualified("ts")),
|
||||
true,
|
||||
true,
|
||||
)];
|
||||
let aggr = |func: Arc<AggregateUDF>, args: Vec<Expr>| {
|
||||
Expr::AggregateFunction(AggregateFunction::new_udf(
|
||||
func,
|
||||
args,
|
||||
false,
|
||||
None,
|
||||
order_by.clone(),
|
||||
None,
|
||||
))
|
||||
};
|
||||
let number = Expr::Column(Column::new_unqualified("number"));
|
||||
|
||||
assert!(!is_all_aggr_exprs_steppable(&[aggr(
|
||||
nth_value_udaf(),
|
||||
vec![number.clone(), lit(2i64)],
|
||||
)]));
|
||||
assert!(is_all_aggr_exprs_steppable(&[aggr(
|
||||
array_agg_udaf(),
|
||||
vec![number.clone()]
|
||||
)]));
|
||||
// WITHIN GROUP (ORDER BY number)
|
||||
assert!(is_all_aggr_exprs_steppable(&[aggr(
|
||||
approx_percentile_cont_udaf(),
|
||||
vec![number, lit(0.5f64)]
|
||||
)]));
|
||||
}
|
||||
|
||||
@@ -221,18 +221,24 @@ impl DfAccumulator for JsonEncodePathAccumulator {
|
||||
)));
|
||||
}
|
||||
|
||||
for state in states {
|
||||
let state = as_struct_array(state)?;
|
||||
let lat_list = as_list_array(state.column(0))?.value(0);
|
||||
let lat_array = as_primitive_array::<Float64Type>(&lat_list)?;
|
||||
let lng_list = as_list_array(state.column(1))?.value(0);
|
||||
let lng_array = as_primitive_array::<Float64Type>(&lng_list)?;
|
||||
let ts_list = as_list_array(state.column(2))?.value(0);
|
||||
let ts_array = as_primitive_array::<Int64Type>(&ts_list)?;
|
||||
let state = as_struct_array(&states[0])?;
|
||||
let lat_lists = as_list_array(state.column(0))?;
|
||||
let lng_lists = as_list_array(state.column(1))?;
|
||||
let ts_lists = as_list_array(state.column(2))?;
|
||||
for row in 0..state.len() {
|
||||
if state.is_null(row) {
|
||||
continue;
|
||||
}
|
||||
let lat_list = lat_lists.value(row);
|
||||
let lng_list = lng_lists.value(row);
|
||||
let ts_list = ts_lists.value(row);
|
||||
|
||||
self.lat.extend(lat_array);
|
||||
self.lng.extend(lng_array);
|
||||
self.timestamp.extend(ts_array);
|
||||
self.lat
|
||||
.extend(as_primitive_array::<Float64Type>(&lat_list)?);
|
||||
self.lng
|
||||
.extend(as_primitive_array::<Float64Type>(&lng_list)?);
|
||||
self.timestamp
|
||||
.extend(as_primitive_array::<Int64Type>(&ts_list)?);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -331,9 +337,9 @@ mod tests {
|
||||
_ => panic!("Expected Struct scalar value"),
|
||||
};
|
||||
|
||||
// Merge state arrays
|
||||
merged.merge_batch(&[state_array1]).unwrap();
|
||||
merged.merge_batch(&[state_array2]).unwrap();
|
||||
// States from different partitions arrive as rows of one batch
|
||||
let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap();
|
||||
merged.merge_batch(&[states]).unwrap();
|
||||
|
||||
// Evaluate merged result
|
||||
let result = merged.evaluate().unwrap();
|
||||
|
||||
@@ -250,18 +250,24 @@ impl DfAccumulator for GeoPathAccumulator {
|
||||
)));
|
||||
}
|
||||
|
||||
for state in states {
|
||||
let state = as_struct_array(state)?;
|
||||
let lat_list = as_list_array(state.column(0))?.value(0);
|
||||
let lat_array = as_primitive_array::<Float64Type>(&lat_list)?;
|
||||
let lng_list = as_list_array(state.column(1))?.value(0);
|
||||
let lng_array = as_primitive_array::<Float64Type>(&lng_list)?;
|
||||
let ts_list = as_list_array(state.column(2))?.value(0);
|
||||
let ts_array = as_primitive_array::<Int64Type>(&ts_list)?;
|
||||
let state = as_struct_array(&states[0])?;
|
||||
let lat_lists = as_list_array(state.column(0))?;
|
||||
let lng_lists = as_list_array(state.column(1))?;
|
||||
let ts_lists = as_list_array(state.column(2))?;
|
||||
for row in 0..state.len() {
|
||||
if state.is_null(row) {
|
||||
continue;
|
||||
}
|
||||
let lat_list = lat_lists.value(row);
|
||||
let lng_list = lng_lists.value(row);
|
||||
let ts_list = ts_lists.value(row);
|
||||
|
||||
self.lat.extend(lat_array);
|
||||
self.lng.extend(lng_array);
|
||||
self.timestamp.extend(ts_array);
|
||||
self.lat
|
||||
.extend(as_primitive_array::<Float64Type>(&lat_list)?);
|
||||
self.lng
|
||||
.extend(as_primitive_array::<Float64Type>(&lng_list)?);
|
||||
self.timestamp
|
||||
.extend(as_primitive_array::<Int64Type>(&ts_list)?);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -403,9 +409,9 @@ mod tests {
|
||||
_ => panic!("Expected Struct scalar value"),
|
||||
};
|
||||
|
||||
// Merge state arrays
|
||||
merged.merge_batch(&[state_array1]).unwrap();
|
||||
merged.merge_batch(&[state_array2]).unwrap();
|
||||
// States from different partitions arrive as rows of one batch
|
||||
let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap();
|
||||
merged.merge_batch(&[states]).unwrap();
|
||||
|
||||
// Evaluate merged result
|
||||
let result = merged.evaluate().unwrap();
|
||||
|
||||
@@ -174,6 +174,15 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Unexpected type {data_type} for column '{column}' of table '{table_name}'"))]
|
||||
UnexpectedColumnType {
|
||||
table_name: String,
|
||||
column: String,
|
||||
data_type: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to collect recordbatch"))]
|
||||
CollectRecordbatch {
|
||||
#[snafu(implicit)]
|
||||
@@ -422,7 +431,9 @@ impl ErrorExt for Error {
|
||||
|
||||
Error::RequestQuery { source, .. } => source.status_code(),
|
||||
|
||||
Error::CacheRequired { .. } => StatusCode::Internal,
|
||||
Error::CacheRequired { .. } | Error::UnexpectedColumnType { .. } => {
|
||||
StatusCode::Internal
|
||||
}
|
||||
|
||||
Error::TableNotFound { .. } => StatusCode::TableNotFound,
|
||||
|
||||
|
||||
@@ -1509,6 +1509,20 @@ impl PrometheusHandler for Instance {
|
||||
.context(ExecuteQuerySnafu)
|
||||
}
|
||||
|
||||
async fn query_metric_names_by_labels(
|
||||
&self,
|
||||
matchers: Vec<Matcher>,
|
||||
schema: &str,
|
||||
start: SystemTime,
|
||||
end: SystemTime,
|
||||
ctx: &QueryContextRef,
|
||||
) -> server_error::Result<Vec<String>> {
|
||||
self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExecuteQuerySnafu)
|
||||
}
|
||||
|
||||
async fn query_label_values(
|
||||
&self,
|
||||
metric: String,
|
||||
|
||||
@@ -519,31 +519,34 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// The PHP instrumentation case: an empty string must be distinguishable
|
||||
/// from any other unparsable value in the reported diagnostics.
|
||||
/// The PHP instrumentation case: `getHeaderLine('Content-Length')` sends ""
|
||||
/// when the header is absent. The span must be kept with a null attribute.
|
||||
#[test]
|
||||
fn test_prepare_trace_column_rewrites_reports_empty_string_value() {
|
||||
let rows = vec![Row {
|
||||
values: vec![Value {
|
||||
value_data: Some(ValueData::StringValue(String::new())),
|
||||
fn test_prepare_trace_column_rewrites_nulls_empty_string_value() {
|
||||
let mut rows = Rows {
|
||||
schema: vec![ColumnSchema {
|
||||
column_name: "span_attributes.http.response.body.size".to_string(),
|
||||
datatype: ColumnDataType::String as i32,
|
||||
..Default::default()
|
||||
}],
|
||||
}];
|
||||
rows: vec![Row {
|
||||
values: vec![Value {
|
||||
value_data: Some(ValueData::StringValue(String::new())),
|
||||
}],
|
||||
}],
|
||||
};
|
||||
let pending_rewrites = vec![PendingTraceColumnRewrite {
|
||||
col_idx: 0,
|
||||
target_type: ColumnDataType::Int64,
|
||||
column_name: "span_attributes.http.response.body.size".to_string(),
|
||||
}];
|
||||
|
||||
let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces")
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.error.to_string().contains(
|
||||
"'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \
|
||||
from String(\"\") to Int64"
|
||||
),
|
||||
"unexpected error message: {}",
|
||||
err.error
|
||||
);
|
||||
prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "opentelemetry_traces")
|
||||
.unwrap()
|
||||
.apply(&mut rows);
|
||||
|
||||
assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32);
|
||||
assert_eq!(rows.rows[0].values[0].value_data, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
@@ -20,20 +21,25 @@ use catalog::information_schema::TABLES;
|
||||
use client::OutputData;
|
||||
use common_catalog::consts::INFORMATION_SCHEMA_NAME;
|
||||
use common_catalog::format_full_table_name;
|
||||
use common_recordbatch::util;
|
||||
use common_recordbatch::{RecordBatch, util};
|
||||
use common_telemetry::tracing;
|
||||
use datafusion_expr::LogicalPlan;
|
||||
use datatypes::arrow::array::{Array, UInt32Array};
|
||||
use futures::StreamExt;
|
||||
use promql_parser::label::{Matcher, Matchers};
|
||||
use query::promql;
|
||||
use query::promql::planner::PromPlanner;
|
||||
use servers::prometheus;
|
||||
use session::context::QueryContextRef;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME;
|
||||
use store_api::storage::TableId;
|
||||
use table::TableRef;
|
||||
|
||||
use crate::error::{
|
||||
CatalogSnafu, CollectRecordbatchSnafu, ExecLogicalPlanSnafu,
|
||||
PrometheusLabelValuesQueryPlanSnafu, PrometheusMetricNamesQueryPlanSnafu, ReadTableSnafu,
|
||||
Result, TableNotFoundSnafu, TableSnafu,
|
||||
Result, TableNotFoundSnafu, TableSnafu, UnexpectedColumnTypeSnafu,
|
||||
};
|
||||
use crate::instance::Instance;
|
||||
|
||||
@@ -117,6 +123,173 @@ impl Instance {
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Handles a metric names query constrained by matchers on ordinary labels.
|
||||
///
|
||||
/// [`Instance::handle_query_metric_names`] answers from table metadata, which
|
||||
/// cannot resolve a label matcher: whether a metric carries `pod="abc"` is a
|
||||
/// property of its data. The metric engine multiplexes the logical tables of
|
||||
/// one physical table into a region holding the union of their label columns
|
||||
/// alongside `__table_id`, so a single distinct scan per physical table
|
||||
/// resolves the matchers for all of its logical tables at once.
|
||||
///
|
||||
/// Only metric engine tables are covered. Tables on other engines share no
|
||||
/// column space to scan, and a scan per table does not scale to the table
|
||||
/// counts this API is expected to answer over.
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub(crate) async fn handle_query_metric_names_by_labels(
|
||||
&self,
|
||||
matchers: Vec<Matcher>,
|
||||
schema: &str,
|
||||
start: SystemTime,
|
||||
end: SystemTime,
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<String>> {
|
||||
let _timer = crate::metrics::PROMQL_QUERY_METRICS_ELAPSED
|
||||
.with_label_values(&[ctx.get_db_string().as_str()])
|
||||
.start_timer();
|
||||
|
||||
let catalog = ctx.current_catalog();
|
||||
let mut table_ids = HashSet::new();
|
||||
for physical in self.physical_metric_tables(catalog, schema, ctx).await? {
|
||||
table_ids.extend(
|
||||
self.scan_matching_table_ids(physical, &matchers, start, end, ctx)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
|
||||
// Batch-resolve only the ids the scan produced. An id dropped between the
|
||||
// scan and here simply has no entry.
|
||||
let table_ids = table_ids.into_iter().collect::<Vec<_>>();
|
||||
let mut names = self
|
||||
.catalog_manager
|
||||
.tables_by_ids(catalog, schema, &table_ids)
|
||||
.await
|
||||
.context(CatalogSnafu)?
|
||||
.into_iter()
|
||||
.map(|table| table.table_info().name.clone())
|
||||
.collect::<Vec<_>>();
|
||||
names.sort_unstable();
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
/// Returns the metric engine physical tables of a schema.
|
||||
///
|
||||
/// Their data regions carry the union of their logical tables' label columns,
|
||||
/// so scanning these covers every metric of the schema.
|
||||
async fn physical_metric_tables(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<TableRef>> {
|
||||
let mut tables = self.catalog_manager.tables(catalog, schema, Some(ctx));
|
||||
let mut physical_tables = Vec::new();
|
||||
|
||||
while let Some(table) = tables.next().await {
|
||||
let table = table.context(CatalogSnafu)?;
|
||||
if table.table_info().is_physical_table() {
|
||||
physical_tables.push(table);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(physical_tables)
|
||||
}
|
||||
|
||||
/// Scans `table` for the distinct values of `column` that match `matchers`
|
||||
/// within the time range.
|
||||
///
|
||||
/// Callers look the table up and decode the batches; the plan between is the
|
||||
/// same whether the projected column is a label or `__table_id`.
|
||||
async fn scan_distinct_column(
|
||||
&self,
|
||||
table: TableRef,
|
||||
matchers: Vec<Matcher>,
|
||||
column: String,
|
||||
start: SystemTime,
|
||||
end: SystemTime,
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<RecordBatch>> {
|
||||
let dataframe = self
|
||||
.query_engine
|
||||
.read_table(table.clone())
|
||||
.with_context(|_| ReadTableSnafu {
|
||||
table_name: table.table_info().full_table_name(),
|
||||
})?;
|
||||
|
||||
let scan_plan = dataframe.into_unoptimized_plan();
|
||||
let conditions = PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema())
|
||||
.context(PrometheusLabelValuesQueryPlanSnafu)?;
|
||||
let logical_plan = promql::label_values::rewrite_label_values_query(
|
||||
table, scan_plan, conditions, column, start, end,
|
||||
)
|
||||
.context(PrometheusLabelValuesQueryPlanSnafu)?;
|
||||
|
||||
let results = self
|
||||
.query_engine
|
||||
.execute(logical_plan, ctx.clone())
|
||||
.await
|
||||
.context(ExecLogicalPlanSnafu)?;
|
||||
|
||||
match results.data {
|
||||
OutputData::Stream(stream) => {
|
||||
util::collect(stream).await.context(CollectRecordbatchSnafu)
|
||||
}
|
||||
OutputData::RecordBatches(rbs) => Ok(rbs.take()),
|
||||
_ => unreachable!("should not happen"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the `__table_id`s of `physical` carrying a row that matches every
|
||||
/// matcher within the time range.
|
||||
async fn scan_matching_table_ids(
|
||||
&self,
|
||||
physical: TableRef,
|
||||
matchers: &[Matcher],
|
||||
start: SystemTime,
|
||||
end: SystemTime,
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<TableId>> {
|
||||
// `__table_id` attributes a row to its logical table, and a physical
|
||||
// table that never took a column from one does not expose it. Skipping
|
||||
// such a table can only miss a metric carrying no labels at all.
|
||||
if physical
|
||||
.schema()
|
||||
.column_schema_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
|
||||
.is_none()
|
||||
{
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let table_name = physical.table_info().full_table_name();
|
||||
let batches = self
|
||||
.scan_distinct_column(
|
||||
physical,
|
||||
matchers.to_vec(),
|
||||
DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
|
||||
start,
|
||||
end,
|
||||
ctx,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut table_ids = Vec::new();
|
||||
for batch in batches {
|
||||
// Only one column in results, ensured by `rewrite_label_values_query`.
|
||||
let column = batch.column(0);
|
||||
let ids = column
|
||||
.as_any()
|
||||
.downcast_ref::<UInt32Array>()
|
||||
.with_context(|| UnexpectedColumnTypeSnafu {
|
||||
table_name: table_name.clone(),
|
||||
column: DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
|
||||
data_type: column.data_type().to_string(),
|
||||
})?;
|
||||
table_ids.extend(ids.iter().flatten());
|
||||
}
|
||||
|
||||
Ok(table_ids)
|
||||
}
|
||||
|
||||
/// Handles label values query request, returns the values.
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub(crate) async fn handle_query_label_values(
|
||||
@@ -150,40 +323,9 @@ impl Instance {
|
||||
.context(TableSnafu);
|
||||
}
|
||||
|
||||
let dataframe = self
|
||||
.query_engine
|
||||
.read_table(table.clone())
|
||||
.with_context(|_| ReadTableSnafu {
|
||||
table_name: full_table_name,
|
||||
})?;
|
||||
|
||||
let scan_plan = dataframe.into_unoptimized_plan();
|
||||
let filter_conditions =
|
||||
PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema())
|
||||
.context(PrometheusLabelValuesQueryPlanSnafu)?;
|
||||
let logical_plan = promql::label_values::rewrite_label_values_query(
|
||||
table,
|
||||
scan_plan,
|
||||
filter_conditions,
|
||||
label_name,
|
||||
start,
|
||||
end,
|
||||
)
|
||||
.context(PrometheusLabelValuesQueryPlanSnafu)?;
|
||||
|
||||
let results = self
|
||||
.query_engine
|
||||
.execute(logical_plan, ctx.clone())
|
||||
.await
|
||||
.context(ExecLogicalPlanSnafu)?;
|
||||
|
||||
let batches = match results.data {
|
||||
OutputData::Stream(stream) => util::collect(stream)
|
||||
.await
|
||||
.context(CollectRecordbatchSnafu)?,
|
||||
OutputData::RecordBatches(rbs) => rbs.take(),
|
||||
_ => unreachable!("should not happen"),
|
||||
};
|
||||
let batches = self
|
||||
.scan_distinct_column(table, matchers, label_name, start, end, ctx)
|
||||
.await?;
|
||||
|
||||
let mut results = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum());
|
||||
for batch in batches {
|
||||
|
||||
@@ -38,7 +38,11 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`;
|
||||
|
||||
- **Heartbeat**: `service/heartbeat.rs` receives datanode/frontend heartbeats and
|
||||
runs the handler chain in `handler/`; responses carry mailbox messages (DDL
|
||||
results, cache invalidation).
|
||||
results, cache invalidation). Node information is refreshed in memory on each
|
||||
heartbeat. Datanode statistics collection and flownode remapping share
|
||||
`handler/utils.rs` for per-node address persistence. Datanode address
|
||||
updates retain the statistics handler's per-node ordering. Address epochs are
|
||||
confirmed only after a successful save; cache invalidation remains best effort.
|
||||
- **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to
|
||||
`common-meta`'s `DdlManager`, executed via `common-procedure`.
|
||||
- **Region migration**: `procedure/region_migration/manager.rs` plus per-step
|
||||
|
||||
@@ -81,6 +81,7 @@ pub mod publish_heartbeat_handler;
|
||||
pub mod region_lease_handler;
|
||||
pub mod remap_flow_peer_handler;
|
||||
pub mod response_header_handler;
|
||||
mod utils;
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_utils;
|
||||
|
||||
@@ -46,30 +46,21 @@ impl HeartbeatHandler for CollectFrontendClusterInfoHandler {
|
||||
ctx: &mut Context,
|
||||
_acc: &mut HeartbeatAccumulator,
|
||||
) -> Result<HandleControl> {
|
||||
let Some((key, peer, info, env_vars)) = extract_base_info(req) else {
|
||||
let Some(update) = NodeInfoUpdate::from_request(req) else {
|
||||
return Ok(HandleControl::Continue);
|
||||
};
|
||||
|
||||
let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref());
|
||||
|
||||
let value = NodeInfo {
|
||||
peer,
|
||||
last_activity_ts: common_time::util::current_time_millis(),
|
||||
status: NodeStatus::Frontend(FrontendStatus {
|
||||
workloads: frontend_workloads,
|
||||
}),
|
||||
version: info.version,
|
||||
git_commit: info.git_commit,
|
||||
start_time_ms: info.start_time_ms,
|
||||
total_cpu_millicores: info.total_cpu_millicores,
|
||||
total_memory_bytes: info.total_memory_bytes,
|
||||
cpu_usage_millicores: info.cpu_usage_millicores,
|
||||
memory_usage_bytes: info.memory_usage_bytes,
|
||||
hostname: info.hostname,
|
||||
env_vars,
|
||||
};
|
||||
|
||||
put_into_memory_store(ctx, key, value).await?;
|
||||
update
|
||||
.save_to_memory(
|
||||
ctx,
|
||||
common_time::util::current_time_millis(),
|
||||
NodeStatus::Frontend(FrontendStatus {
|
||||
workloads: frontend_workloads,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(HandleControl::Continue)
|
||||
}
|
||||
@@ -89,29 +80,20 @@ impl HeartbeatHandler for CollectFlownodeClusterInfoHandler {
|
||||
ctx: &mut Context,
|
||||
_acc: &mut HeartbeatAccumulator,
|
||||
) -> Result<HandleControl> {
|
||||
let Some((key, peer, info, env_vars)) = extract_base_info(req) else {
|
||||
let Some(update) = NodeInfoUpdate::from_request(req) else {
|
||||
return Ok(HandleControl::Continue);
|
||||
};
|
||||
let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref());
|
||||
|
||||
let value = NodeInfo {
|
||||
peer,
|
||||
last_activity_ts: common_time::util::current_time_millis(),
|
||||
status: NodeStatus::Flownode(FlownodeStatus {
|
||||
workloads: flownode_workloads,
|
||||
}),
|
||||
version: info.version,
|
||||
git_commit: info.git_commit,
|
||||
start_time_ms: info.start_time_ms,
|
||||
total_cpu_millicores: info.total_cpu_millicores,
|
||||
total_memory_bytes: info.total_memory_bytes,
|
||||
cpu_usage_millicores: info.cpu_usage_millicores,
|
||||
memory_usage_bytes: info.memory_usage_bytes,
|
||||
hostname: info.hostname,
|
||||
env_vars,
|
||||
};
|
||||
|
||||
put_into_memory_store(ctx, key, value).await?;
|
||||
update
|
||||
.save_to_memory(
|
||||
ctx,
|
||||
common_time::util::current_time_millis(),
|
||||
NodeStatus::Flownode(FlownodeStatus {
|
||||
workloads: flownode_workloads,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(HandleControl::Continue)
|
||||
}
|
||||
@@ -132,7 +114,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler {
|
||||
ctx: &mut Context,
|
||||
acc: &mut HeartbeatAccumulator,
|
||||
) -> Result<HandleControl> {
|
||||
let Some((key, peer, info, env_vars)) = extract_base_info(req) else {
|
||||
let Some(update) = NodeInfoUpdate::from_request(req) else {
|
||||
return Ok(HandleControl::Continue);
|
||||
};
|
||||
|
||||
@@ -147,16 +129,64 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler {
|
||||
.count();
|
||||
let follower_regions = stat.region_stats.len() - leader_regions;
|
||||
|
||||
let value = NodeInfo {
|
||||
update
|
||||
.save_to_memory(
|
||||
ctx,
|
||||
stat.timestamp_millis,
|
||||
NodeStatus::Datanode(DatanodeStatus {
|
||||
rcus: stat.rcus,
|
||||
wcus: stat.wcus,
|
||||
leader_regions,
|
||||
follower_regions,
|
||||
workloads: stat.datanode_workloads.clone(),
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(HandleControl::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
/// Common fields collected for one heartbeat's in-memory node information
|
||||
/// update. Persistent routing addresses have a separate lifecycle.
|
||||
struct NodeInfoUpdate {
|
||||
key: NodeInfoKey,
|
||||
peer: Peer,
|
||||
info: PbNodeInfo,
|
||||
env_vars: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl NodeInfoUpdate {
|
||||
fn from_request(request: &HeartbeatRequest) -> Option<Self> {
|
||||
let key = NodeInfoKey::new(request)?;
|
||||
let peer = request.peer.clone()?;
|
||||
let info = request.info.clone()?;
|
||||
let env_vars = EnvVars::from_extensions(&request.extensions)
|
||||
.inspect_err(|e| {
|
||||
warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer);
|
||||
})
|
||||
.unwrap_or_default()
|
||||
.map(|e| e.vars)
|
||||
.unwrap_or_default();
|
||||
Some(Self {
|
||||
key,
|
||||
peer,
|
||||
last_activity_ts: stat.timestamp_millis,
|
||||
status: NodeStatus::Datanode(DatanodeStatus {
|
||||
rcus: stat.rcus,
|
||||
wcus: stat.wcus,
|
||||
leader_regions,
|
||||
follower_regions,
|
||||
workloads: stat.datanode_workloads.clone(),
|
||||
}),
|
||||
info,
|
||||
env_vars,
|
||||
})
|
||||
}
|
||||
|
||||
async fn save_to_memory(
|
||||
self,
|
||||
ctx: &Context,
|
||||
last_activity_ts: i64,
|
||||
status: NodeStatus,
|
||||
) -> Result<()> {
|
||||
let info = self.info;
|
||||
let value = NodeInfo {
|
||||
peer: self.peer,
|
||||
last_activity_ts,
|
||||
status,
|
||||
version: info.version,
|
||||
git_commit: info.git_commit,
|
||||
start_time_ms: info.start_time_ms,
|
||||
@@ -165,53 +195,88 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler {
|
||||
cpu_usage_millicores: info.cpu_usage_millicores,
|
||||
memory_usage_bytes: info.memory_usage_bytes,
|
||||
hostname: info.hostname,
|
||||
env_vars,
|
||||
env_vars: self.env_vars,
|
||||
};
|
||||
|
||||
put_into_memory_store(ctx, key, value).await?;
|
||||
|
||||
Ok(HandleControl::Continue)
|
||||
let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?;
|
||||
ctx.in_memory
|
||||
.put(PutRequest {
|
||||
key: (&self.key).into(),
|
||||
value,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.context(SaveClusterInfoSnafu)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_base_info(
|
||||
request: &HeartbeatRequest,
|
||||
) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap<String, String>)> {
|
||||
let HeartbeatRequest { peer, info, .. } = request;
|
||||
let key = NodeInfoKey::new(request)?;
|
||||
let Some(peer) = &peer else {
|
||||
return None;
|
||||
};
|
||||
let Some(info) = &info else {
|
||||
return None;
|
||||
};
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use api::v1::meta::RequestHeader;
|
||||
use common_meta::datanode::Stat;
|
||||
|
||||
let env_vars = EnvVars::from_extensions(&request.extensions)
|
||||
.inspect_err(|e| {
|
||||
warn!(e;
|
||||
"Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer
|
||||
use super::*;
|
||||
use crate::handler::test_utils::TestEnv;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_node_info_refreshes_without_epoch_change() {
|
||||
let mut ctx = TestEnv::new().ctx();
|
||||
let handlers: Vec<(Role, Box<dyn HeartbeatHandler>)> = vec![
|
||||
(Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)),
|
||||
(Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)),
|
||||
(Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)),
|
||||
];
|
||||
for (role, handler) in handlers {
|
||||
let mut req = HeartbeatRequest {
|
||||
header: Some(RequestHeader {
|
||||
role: role as i32,
|
||||
..Default::default()
|
||||
}),
|
||||
peer: Some(Peer {
|
||||
id: 7,
|
||||
addr: "node:3001".into(),
|
||||
}),
|
||||
info: Some(PbNodeInfo {
|
||||
version: "v1".into(),
|
||||
hostname: "host".into(),
|
||||
..Default::default()
|
||||
}),
|
||||
node_epoch: 10,
|
||||
..Default::default()
|
||||
};
|
||||
let mut acc = HeartbeatAccumulator {
|
||||
stat: Some(Stat {
|
||||
timestamp_millis: 123,
|
||||
rcus: 5,
|
||||
wcus: 6,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
|
||||
let key: Vec<u8> = (&NodeInfoKey::new(&req).unwrap()).into();
|
||||
let first =
|
||||
NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap();
|
||||
assert_eq!(first.peer, req.peer.clone().unwrap());
|
||||
assert_eq!(first.version, "v1");
|
||||
assert_eq!(first.hostname, "host");
|
||||
req.info.as_mut().unwrap().version = "v2".into();
|
||||
req.info.as_mut().unwrap().memory_usage_bytes = 42;
|
||||
acc.stat.as_mut().unwrap().timestamp_millis = 456;
|
||||
handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
|
||||
let second =
|
||||
NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap();
|
||||
assert_eq!(second.version, "v2");
|
||||
assert_eq!(second.memory_usage_bytes, 42);
|
||||
assert_eq!(
|
||||
serde_json::to_value(&second.status).unwrap(),
|
||||
serde_json::to_value(&first.status).unwrap()
|
||||
);
|
||||
})
|
||||
.unwrap_or_default()
|
||||
.map(|e| e.vars)
|
||||
.unwrap_or_default();
|
||||
|
||||
Some((key, peer.clone(), info.clone(), env_vars))
|
||||
}
|
||||
|
||||
async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> {
|
||||
let key = (&key).into();
|
||||
let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?;
|
||||
let put_req = PutRequest {
|
||||
key,
|
||||
value,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
ctx.in_memory
|
||||
.put(put_req)
|
||||
.await
|
||||
.context(SaveClusterInfoSnafu)?;
|
||||
|
||||
Ok(())
|
||||
if role == Role::Datanode {
|
||||
assert_eq!(second.last_activity_ts, 456);
|
||||
}
|
||||
// Node information must not be persisted to the configured backend.
|
||||
assert!(ctx.kv_backend.get(&key).await.unwrap().is_none());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,19 +15,19 @@
|
||||
use std::cmp::Ordering;
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::meta::{HeartbeatRequest, Role};
|
||||
use api::v1::meta::{HeartbeatRequest, Peer, Role};
|
||||
use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat};
|
||||
use common_meta::instruction::CacheIdent;
|
||||
use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue};
|
||||
use common_meta::key::{MetadataKey, MetadataValue};
|
||||
use common_meta::peer::Peer;
|
||||
use common_meta::key::MetadataKey;
|
||||
use common_meta::key::node_address::NodeAddressKey;
|
||||
use common_meta::rpc::store::PutRequest;
|
||||
use common_telemetry::{error, info, warn};
|
||||
use common_telemetry::{error, warn};
|
||||
use dashmap::DashMap;
|
||||
use snafu::ResultExt;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::error::{self, Result};
|
||||
use crate::handler::utils::{self, NodeAddressUpdater};
|
||||
use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
|
||||
use crate::metasrv::Context;
|
||||
|
||||
@@ -74,6 +74,7 @@ const DEFAULT_FLUSH_STATS_FACTOR: usize = 3;
|
||||
pub struct CollectStatsHandler {
|
||||
stats_cache: DashMap<DatanodeStatKey, Arc<Mutex<EpochStats>>>,
|
||||
flush_stats_factor: usize,
|
||||
address_updater: NodeAddressUpdater,
|
||||
}
|
||||
|
||||
impl Default for CollectStatsHandler {
|
||||
@@ -87,6 +88,7 @@ impl CollectStatsHandler {
|
||||
Self {
|
||||
flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR),
|
||||
stats_cache: DashMap::default(),
|
||||
address_updater: NodeAddressUpdater::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -119,6 +121,7 @@ impl HeartbeatHandler for CollectStatsHandler {
|
||||
|
||||
let key: Vec<u8> = key.into();
|
||||
|
||||
let current_epoch = current_stat.node_epoch;
|
||||
let refresh = if let Some(epoch) = epoch_stats.epoch() {
|
||||
match current_stat.node_epoch.cmp(&epoch) {
|
||||
Ordering::Greater => {
|
||||
@@ -146,11 +149,13 @@ impl HeartbeatHandler for CollectStatsHandler {
|
||||
true
|
||||
};
|
||||
|
||||
// Need to refresh the [datanode -> address] mapping
|
||||
if refresh {
|
||||
// Safety: `epoch_stats.stats` is not empty
|
||||
let last = epoch_stats.stats.last().unwrap();
|
||||
rewrite_node_address(ctx, last).await;
|
||||
// Keep address updates and statistics in the original per-node order.
|
||||
// Address acknowledgement is independent of the statistics epoch so a
|
||||
// failed save can be retried by the next heartbeat of the same epoch.
|
||||
if epoch_stats.epoch() == Some(current_epoch)
|
||||
&& let Some(stat) = epoch_stats.stats.last()
|
||||
{
|
||||
update_datanode_address(&self.address_updater, ctx, stat).await;
|
||||
}
|
||||
|
||||
if !refresh && epoch_stats.len() < self.flush_stats_factor {
|
||||
@@ -178,48 +183,26 @@ impl HeartbeatHandler for CollectStatsHandler {
|
||||
}
|
||||
}
|
||||
|
||||
async fn rewrite_node_address(ctx: &Context, stat: &Stat) {
|
||||
let peer = Peer {
|
||||
id: stat.id,
|
||||
addr: stat.addr.clone(),
|
||||
};
|
||||
let key = NodeAddressKey::with_datanode(peer.id).to_bytes();
|
||||
if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() {
|
||||
let put = PutRequest {
|
||||
key,
|
||||
value,
|
||||
prev_kv: false,
|
||||
};
|
||||
|
||||
match ctx.leader_cached_kv_backend.put(put).await {
|
||||
Ok(_) => {
|
||||
info!(
|
||||
"Successfully updated datanode `NodeAddressValue`: {:?}",
|
||||
peer
|
||||
);
|
||||
// broadcast invalidating cache
|
||||
let cache_idents = stat
|
||||
.table_ids()
|
||||
.into_iter()
|
||||
.map(CacheIdent::TableId)
|
||||
.collect::<Vec<_>>();
|
||||
if let Err(e) = ctx
|
||||
.cache_invalidator
|
||||
.invalidate(&Default::default(), &cache_idents)
|
||||
.await
|
||||
{
|
||||
error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!(e; "Failed to update datanode `NodeAddressValue`: {:?}", peer);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
warn!(
|
||||
"Failed to serialize datanode `NodeAddressValue`: {:?}",
|
||||
peer
|
||||
);
|
||||
/// Applies datanode-specific address and table-cache updates. The caller keeps
|
||||
/// its statistics lock so notifications and statistics retain their ordering.
|
||||
async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) {
|
||||
let updated = updater
|
||||
.update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async {
|
||||
let peer = Peer { id: stat.id, addr: stat.addr.clone() };
|
||||
let key = NodeAddressKey::with_datanode(stat.id).to_bytes();
|
||||
utils::save_node_address(ctx, key, peer).await
|
||||
})
|
||||
.await
|
||||
.inspect_err(|e| {
|
||||
error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr);
|
||||
});
|
||||
if let Ok(true) = updated {
|
||||
let cache_idents = stat
|
||||
.table_ids()
|
||||
.into_iter()
|
||||
.map(CacheIdent::TableId)
|
||||
.collect::<Vec<_>>();
|
||||
utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -553,4 +536,127 @@ mod tests {
|
||||
handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_persistence_failure_retries_but_notification_failure_does_not() {
|
||||
use api::v1::meta::HeartbeatRequest;
|
||||
use common_meta::key::MetadataValue;
|
||||
use common_meta::key::node_address::NodeAddressValue;
|
||||
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
|
||||
use common_meta::rpc::store::PutResponse;
|
||||
|
||||
use crate::handler::collect_stats_handler::CollectStatsHandler;
|
||||
use crate::handler::test_utils::FailingCacheInvalidator;
|
||||
use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
|
||||
use crate::service::store::cached_kv::LeaderCachedKvBackend;
|
||||
let mut ctx = TestEnv::new().ctx();
|
||||
let puts = Arc::new(std::sync::Mutex::new(Vec::<PutRequest>::new()));
|
||||
let backend = MockKvBackendBuilder::default()
|
||||
.put_fn(Arc::new({
|
||||
let puts = puts.clone();
|
||||
move |req: PutRequest| {
|
||||
let mut puts = puts.lock().unwrap();
|
||||
puts.push(req);
|
||||
if puts.len() == 1 {
|
||||
return common_meta::error::UnexpectedSnafu {
|
||||
err_msg: "injected persistence failure",
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
Ok(PutResponse::default())
|
||||
}
|
||||
}) as _)
|
||||
.build()
|
||||
.unwrap();
|
||||
ctx.leader_cached_kv_backend =
|
||||
Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend)));
|
||||
let invalidator = Arc::new(FailingCacheInvalidator::default());
|
||||
ctx.cache_invalidator = invalidator.clone();
|
||||
let handler = CollectStatsHandler::default();
|
||||
let mut stat = Stat {
|
||||
id: 101,
|
||||
addr: "dn-original".into(),
|
||||
node_epoch: 7,
|
||||
..Default::default()
|
||||
};
|
||||
let req = HeartbeatRequest::default();
|
||||
|
||||
// A backend error must not interrupt subsequent heartbeat handlers.
|
||||
assert!(matches!(
|
||||
handler
|
||||
.handle(
|
||||
&req,
|
||||
&mut ctx,
|
||||
&mut HeartbeatAccumulator {
|
||||
stat: Some(stat.clone()),
|
||||
..Default::default()
|
||||
}
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
HandleControl::Continue
|
||||
));
|
||||
assert_eq!(puts.lock().unwrap().len(), 1);
|
||||
assert_eq!(invalidator.attempts(), 0);
|
||||
|
||||
let stats_key: Vec<u8> = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into();
|
||||
assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some());
|
||||
|
||||
// Same epoch retries persistence using the current heartbeat's address.
|
||||
stat.addr = "dn-retry".into();
|
||||
handler
|
||||
.handle(
|
||||
&req,
|
||||
&mut ctx,
|
||||
&mut HeartbeatAccumulator {
|
||||
stat: Some(stat.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
{
|
||||
let puts = puts.lock().unwrap();
|
||||
assert_eq!(puts.len(), 2);
|
||||
assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes());
|
||||
assert_eq!(
|
||||
NodeAddressValue::try_from_raw_value(&puts[1].value)
|
||||
.unwrap()
|
||||
.peer
|
||||
.addr,
|
||||
"dn-retry"
|
||||
);
|
||||
}
|
||||
assert_eq!(invalidator.attempts(), 1);
|
||||
|
||||
// Notification remains best-effort: its failure neither rewrites nor retries.
|
||||
handler
|
||||
.handle(
|
||||
&req,
|
||||
&mut ctx,
|
||||
&mut HeartbeatAccumulator {
|
||||
stat: Some(stat.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(puts.lock().unwrap().len(), 2);
|
||||
assert_eq!(invalidator.attempts(), 1);
|
||||
|
||||
stat.node_epoch = 8;
|
||||
handler
|
||||
.handle(
|
||||
&req,
|
||||
&mut ctx,
|
||||
&mut HeartbeatAccumulator {
|
||||
stat: Some(stat.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(puts.lock().unwrap().len(), 3);
|
||||
assert_eq!(invalidator.attempts(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,22 +12,20 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use api::v1::meta::{HeartbeatRequest, Peer, Role};
|
||||
use api::v1::meta::{HeartbeatRequest, Role};
|
||||
use common_meta::instruction::CacheIdent;
|
||||
use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue};
|
||||
use common_meta::key::{MetadataKey, MetadataValue};
|
||||
use common_meta::rpc::store::PutRequest;
|
||||
use common_telemetry::{error, info, warn};
|
||||
use dashmap::DashMap;
|
||||
use common_meta::key::MetadataKey;
|
||||
use common_meta::key::node_address::NodeAddressKey;
|
||||
use common_telemetry::error;
|
||||
|
||||
use crate::Result;
|
||||
use crate::handler::utils::{self, NodeAddressUpdater};
|
||||
use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
|
||||
use crate::metasrv::Context;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct RemapFlowPeerHandler {
|
||||
/// flow_node_id -> epoch
|
||||
epoch_cache: DashMap<u64, u64>,
|
||||
address_updater: NodeAddressUpdater,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
@@ -45,61 +43,85 @@ impl HeartbeatHandler for RemapFlowPeerHandler {
|
||||
let Some(peer) = req.peer.as_ref() else {
|
||||
return Ok(HandleControl::Continue);
|
||||
};
|
||||
|
||||
let current_epoch = req.node_epoch;
|
||||
let flow_node_id = peer.id;
|
||||
|
||||
let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) {
|
||||
if current_epoch > *epoch.value() {
|
||||
*epoch.value_mut() = current_epoch;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
self.epoch_cache.insert(flow_node_id, current_epoch);
|
||||
true
|
||||
};
|
||||
|
||||
if refresh {
|
||||
rewrite_node_address(ctx, peer).await;
|
||||
let updated = self
|
||||
.address_updater
|
||||
.update_if_needed((Role::Flownode, peer.id), req.node_epoch, || {
|
||||
let key = NodeAddressKey::with_flownode(peer.id).to_bytes();
|
||||
utils::save_node_address(ctx, key, peer.clone())
|
||||
})
|
||||
.await
|
||||
.inspect_err(|e| {
|
||||
// Address errors must not prevent later heartbeat handlers from running.
|
||||
error!(e; "Failed to update flownode address, peer: {:?}", peer);
|
||||
});
|
||||
if let Ok(true) = updated {
|
||||
let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)];
|
||||
utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await;
|
||||
}
|
||||
|
||||
Ok(HandleControl::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) {
|
||||
let key = NodeAddressKey::with_flownode(peer.id).to_bytes();
|
||||
if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() {
|
||||
let put = PutRequest {
|
||||
key,
|
||||
value,
|
||||
prev_kv: false,
|
||||
};
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
match ctx.leader_cached_kv_backend.put(put).await {
|
||||
Ok(_) => {
|
||||
info!("Successfully updated flow `NodeAddressValue`: {:?}", peer);
|
||||
// broadcast invalidating cache to all frontends
|
||||
let cache_idents = vec![CacheIdent::FlowNodeAddressChange(peer.id)];
|
||||
info!(
|
||||
"Invalidate flow node cache for new address with cache idents: {:?}",
|
||||
cache_idents
|
||||
);
|
||||
if let Err(e) = ctx
|
||||
.cache_invalidator
|
||||
.invalidate(&Default::default(), &cache_idents)
|
||||
.await
|
||||
{
|
||||
error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer);
|
||||
use api::v1::meta::Peer;
|
||||
|
||||
use super::*;
|
||||
use crate::handler::test_utils::TestEnv;
|
||||
#[tokio::test]
|
||||
async fn test_notification_failure_keeps_persisted_epoch() {
|
||||
use api::v1::meta::HeartbeatRequest;
|
||||
use common_meta::key::MetadataValue;
|
||||
use common_meta::key::node_address::NodeAddressValue;
|
||||
|
||||
use crate::handler::test_utils::FailingCacheInvalidator;
|
||||
use crate::handler::{HeartbeatAccumulator, HeartbeatHandler};
|
||||
let mut ctx = TestEnv::new().ctx();
|
||||
let invalidator = Arc::new(FailingCacheInvalidator::default());
|
||||
ctx.cache_invalidator = invalidator.clone();
|
||||
let handler = RemapFlowPeerHandler::default();
|
||||
let mut req = HeartbeatRequest {
|
||||
peer: Some(Peer {
|
||||
id: 1,
|
||||
addr: "first".into(),
|
||||
}),
|
||||
node_epoch: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let key = NodeAddressKey::with_flownode(1).to_bytes();
|
||||
handler
|
||||
.handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
|
||||
.await
|
||||
.unwrap();
|
||||
req.peer.as_mut().unwrap().addr = "same-epoch".into();
|
||||
handler
|
||||
.handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(invalidator.attempts(), 1);
|
||||
let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
NodeAddressValue::try_from_raw_value(&kv.value)
|
||||
.unwrap()
|
||||
.peer
|
||||
.addr,
|
||||
"first"
|
||||
);
|
||||
req.node_epoch = 2;
|
||||
handler
|
||||
.handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(invalidator.attempts(), 2);
|
||||
let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
NodeAddressValue::try_from_raw_value(&kv.value)
|
||||
.unwrap()
|
||||
.peer
|
||||
.addr,
|
||||
"same-epoch"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,9 +13,13 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator};
|
||||
use common_meta::cache_invalidator::{
|
||||
CacheInvalidator, CacheInvalidatorRef, Context as InvalidationContext, DummyCacheInvalidator,
|
||||
};
|
||||
use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL;
|
||||
use common_meta::instruction::CacheIdent;
|
||||
use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef};
|
||||
@@ -97,3 +101,34 @@ impl TestEnv {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Counts failed notification attempts without requiring heartbeat connections.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct FailingCacheInvalidator {
|
||||
attempts: AtomicUsize,
|
||||
}
|
||||
|
||||
impl FailingCacheInvalidator {
|
||||
pub(crate) fn attempts(&self) -> usize {
|
||||
self.attempts.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl CacheInvalidator for FailingCacheInvalidator {
|
||||
async fn invalidate(
|
||||
&self,
|
||||
_: &InvalidationContext,
|
||||
_: &[CacheIdent],
|
||||
) -> common_meta::error::Result<()> {
|
||||
self.attempts.fetch_add(1, Ordering::Relaxed);
|
||||
common_meta::error::UnexpectedSnafu {
|
||||
err_msg: "injected notification failure",
|
||||
}
|
||||
.fail()
|
||||
}
|
||||
|
||||
fn invalidate_all(&self) -> common_meta::error::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,281 @@
|
||||
// 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::future::Future;
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::meta::{Peer, Role};
|
||||
use common_meta::instruction::CacheIdent;
|
||||
use common_meta::key::MetadataValue;
|
||||
use common_meta::key::node_address::NodeAddressValue;
|
||||
use common_meta::rpc::store::PutRequest;
|
||||
use common_telemetry::{error, info};
|
||||
use dashmap::DashMap;
|
||||
use snafu::ResultExt;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::Result;
|
||||
use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu};
|
||||
use crate::metasrv::Context;
|
||||
|
||||
/// Serializes address changes per node and acknowledges only successful updates.
|
||||
/// Node information refreshed on every heartbeat has a separate lifecycle.
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct NodeAddressUpdater {
|
||||
states: DashMap<(Role, u64), Arc<Mutex<NodeAddressEpochs>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct NodeAddressEpochs {
|
||||
/// Rejects stale heartbeats even after a newer epoch failed to save.
|
||||
seen_epoch: Option<u64>,
|
||||
/// Suppresses retries only after persistence succeeds.
|
||||
saved_epoch: Option<u64>,
|
||||
}
|
||||
|
||||
impl NodeAddressUpdater {
|
||||
/// Saves a new or previously failed epoch. Returns true only when this call
|
||||
/// saved the address, or false when the heartbeat is stale or already saved.
|
||||
/// Cache invalidation is separate and does not affect acknowledgement.
|
||||
pub(crate) async fn update_if_needed<F, Fut>(
|
||||
&self,
|
||||
node: (Role, u64),
|
||||
epoch: u64,
|
||||
save: F,
|
||||
) -> Result<bool>
|
||||
where
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<()>>,
|
||||
{
|
||||
// Never retain a DashMap guard across await, even while waiting for a
|
||||
// different node: its entry may live in the same shard.
|
||||
let state = {
|
||||
let entry = self.states.entry(node).or_default();
|
||||
Arc::clone(entry.value())
|
||||
};
|
||||
let mut state = state.lock().await;
|
||||
if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// Seeing an epoch rejects stale heartbeats, but only a successful save
|
||||
// suppresses retries of that same epoch. Cancellation also leaves it retryable.
|
||||
state.seen_epoch = Some(epoch);
|
||||
save().await?;
|
||||
state.saved_epoch = Some(epoch);
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
/// Persists an address through the leader cache.
|
||||
pub(crate) async fn save_node_address(ctx: &Context, key: Vec<u8>, peer: Peer) -> Result<()> {
|
||||
let address = NodeAddressValue::new(peer);
|
||||
let value = address
|
||||
.try_as_raw_value()
|
||||
.context(InvalidNodeInfoFormatSnafu)?;
|
||||
ctx.leader_cached_kv_backend
|
||||
.put(PutRequest {
|
||||
key,
|
||||
value,
|
||||
prev_kv: false,
|
||||
})
|
||||
.await
|
||||
.context(KvBackendSnafu)?;
|
||||
info!("Successfully updated node address: {:?}", address.peer);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Preserves the existing best-effort notification behavior. Notification
|
||||
/// failures do not undo a persisted address or cause same-epoch rewrites.
|
||||
pub(crate) async fn invalidate_address_caches(
|
||||
ctx: &Context,
|
||||
node_id: u64,
|
||||
cache_idents: &[CacheIdent],
|
||||
) {
|
||||
if let Err(e) = ctx
|
||||
.cache_invalidator
|
||||
.invalidate(&Default::default(), cache_idents)
|
||||
.await
|
||||
{
|
||||
error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::Notify;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() {
|
||||
let updater = NodeAddressUpdater::default();
|
||||
let key = (Role::Datanode, 1);
|
||||
assert!(
|
||||
updater
|
||||
.update_if_needed(key, 2, || async {
|
||||
crate::error::UnexpectedSnafu {
|
||||
violated: "injected save failure",
|
||||
}
|
||||
.fail()
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
!updater
|
||||
.update_if_needed(key, 1, || async {
|
||||
panic!("stale heartbeat must not save");
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert!(
|
||||
updater
|
||||
.update_if_needed(key, 2, || async { Ok(()) })
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert!(
|
||||
!updater
|
||||
.update_if_needed(key, 2, || async {
|
||||
panic!("successful epoch must not save again");
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert!(
|
||||
updater
|
||||
.update_if_needed(key, 3, || async { Ok(()) })
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancelled_save_can_be_retried() {
|
||||
let updater = NodeAddressUpdater::default();
|
||||
let key = (Role::Datanode, 1);
|
||||
let entered = Notify::new();
|
||||
{
|
||||
let first = updater.update_if_needed(key, 0, || async {
|
||||
entered.notify_one();
|
||||
std::future::pending::<Result<()>>().await
|
||||
});
|
||||
tokio::pin!(first);
|
||||
tokio::select! {
|
||||
_ = entered.notified() => {},
|
||||
_ = &mut first => panic!("save must remain pending"),
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
timeout(
|
||||
Duration::from_secs(2),
|
||||
updater.update_if_needed(key, 0, || async { Ok(()) })
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_same_node_serialized_other_nodes_progress() {
|
||||
// A watchdog outside the runtime detects synchronous DashMap deadlocks;
|
||||
// a Tokio timeout alone cannot fire if the only worker is blocked.
|
||||
let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1);
|
||||
let release = Arc::new(Notify::new());
|
||||
let finish = Arc::new(Notify::new());
|
||||
let worker = std::thread::spawn({
|
||||
let release = release.clone();
|
||||
let finish = finish.clone();
|
||||
move || {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_time()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async move {
|
||||
let updater = Arc::new(NodeAddressUpdater::default());
|
||||
let entered = Arc::new(Notify::new());
|
||||
|
||||
let first = tokio::spawn({
|
||||
let updater = updater.clone();
|
||||
let entered = entered.clone();
|
||||
let release = release.clone();
|
||||
async move {
|
||||
updater
|
||||
.update_if_needed((Role::Datanode, 1), 1, || async {
|
||||
entered.notify_one();
|
||||
release.notified().await;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
});
|
||||
entered.notified().await;
|
||||
let second = tokio::spawn({
|
||||
let updater = updater.clone();
|
||||
async move {
|
||||
updater
|
||||
.update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) })
|
||||
.await
|
||||
}
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
assert!(!second.is_finished());
|
||||
// Exercise shared DashMap shards on a single-thread runtime while the
|
||||
// first node is suspended in external I/O.
|
||||
timeout(Duration::from_secs(2), async {
|
||||
for id in 0_u64..128 {
|
||||
assert!(
|
||||
updater
|
||||
.update_if_needed((Role::Datanode, id + 2), 1, || async {
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
progress_tx.send(()).unwrap();
|
||||
finish.notified().await;
|
||||
timeout(Duration::from_secs(2), async {
|
||||
assert!(first.await.unwrap().unwrap());
|
||||
assert!(second.await.unwrap().unwrap());
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
!updater
|
||||
.update_if_needed((Role::Datanode, 1), 1, || async {
|
||||
panic!("old epoch must not overwrite a newer one");
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
});
|
||||
}
|
||||
});
|
||||
let progress = progress_rx.recv_timeout(Duration::from_secs(5));
|
||||
release.notify_one();
|
||||
finish.notify_one();
|
||||
progress.expect("other node updates must progress while one node save is blocked");
|
||||
worker.join().unwrap();
|
||||
}
|
||||
}
|
||||
@@ -510,17 +510,14 @@ impl RegionMigrationManager {
|
||||
ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
|
||||
let procedure_id = procedure_with_id.id;
|
||||
info!("Starting region migration procedure {procedure_id} for {task}");
|
||||
let procedure_manager = self.procedure_manager.clone();
|
||||
let mut watcher = self
|
||||
.procedure_manager
|
||||
.submit(procedure_with_id)
|
||||
.await
|
||||
.context(error::SubmitProcedureSnafu)?;
|
||||
let num_region = task.region_ids.len();
|
||||
|
||||
common_runtime::spawn_global(async move {
|
||||
let watcher = &mut match procedure_manager.submit(procedure_with_id).await {
|
||||
Ok(watcher) => watcher,
|
||||
Err(e) => {
|
||||
error!(e; "Failed to submit region migration procedure {procedure_id} for {task}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
METRIC_META_REGION_MIGRATION_DATANODES
|
||||
.with_label_values(&["src", &task.from_peer.id.to_string()])
|
||||
.inc_by(num_region as u64);
|
||||
@@ -528,7 +525,7 @@ impl RegionMigrationManager {
|
||||
.with_label_values(&["desc", &task.to_peer.id.to_string()])
|
||||
.inc_by(num_region as u64);
|
||||
|
||||
if let Err(e) = watcher::wait(watcher).await {
|
||||
if let Err(e) = watcher::wait(&mut watcher).await {
|
||||
error!(e; "Failed to wait region migration procedure {procedure_id} for {task}");
|
||||
METRIC_META_REGION_MIGRATION_FAIL.inc();
|
||||
return;
|
||||
@@ -613,15 +610,12 @@ impl RegionMigrationManager {
|
||||
ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
|
||||
let procedure_id = procedure_with_id.id;
|
||||
info!("Starting region migration procedure {procedure_id} for {task}");
|
||||
let procedure_manager = self.procedure_manager.clone();
|
||||
let mut watcher = self
|
||||
.procedure_manager
|
||||
.submit(procedure_with_id)
|
||||
.await
|
||||
.context(error::SubmitProcedureSnafu)?;
|
||||
common_runtime::spawn_global(async move {
|
||||
let watcher = &mut match procedure_manager.submit(procedure_with_id).await {
|
||||
Ok(watcher) => watcher,
|
||||
Err(e) => {
|
||||
error!(e; "Failed to submit region migration procedure {procedure_id} for {task}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
METRIC_META_REGION_MIGRATION_DATANODES
|
||||
.with_label_values(&["src", &task.from_peer.id.to_string()])
|
||||
.inc();
|
||||
@@ -629,7 +623,7 @@ impl RegionMigrationManager {
|
||||
.with_label_values(&["desc", &task.to_peer.id.to_string()])
|
||||
.inc();
|
||||
|
||||
if let Err(e) = watcher::wait(watcher).await {
|
||||
if let Err(e) = watcher::wait(&mut watcher).await {
|
||||
error!(e; "Failed to wait region migration procedure {procedure_id} for {task}");
|
||||
METRIC_META_REGION_MIGRATION_FAIL.inc();
|
||||
return;
|
||||
@@ -649,10 +643,148 @@ mod test {
|
||||
use common_meta::key::table_route::LogicalTableRouteValue;
|
||||
use common_meta::key::test_utils::new_test_table_info;
|
||||
use common_meta::rpc::router::Region;
|
||||
use common_meta::state_store::KvStateStore;
|
||||
use common_procedure::ProcedureManager;
|
||||
use common_procedure::local::{LocalManager, ManagerConfig, PauseAware};
|
||||
use common_procedure::test_util::InMemoryPoisonStore;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use super::*;
|
||||
use crate::procedure::region_migration::test_util::TestingEnv;
|
||||
|
||||
#[derive(Default)]
|
||||
struct SubmissionGate {
|
||||
entered: Notify,
|
||||
release: Notify,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl PauseAware for SubmissionGate {
|
||||
async fn is_paused(&self) -> std::result::Result<bool, common_error::ext::BoxedError> {
|
||||
self.entered.notify_one();
|
||||
self.release.notified().await;
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
async fn check_submission_registration(batch: bool) {
|
||||
let env = TestingEnv::new();
|
||||
let gate = Arc::new(SubmissionGate::default());
|
||||
let procedure_manager = Arc::new(LocalManager::new(
|
||||
ManagerConfig::default(),
|
||||
Arc::new(KvStateStore::new(env.kv_backend())),
|
||||
Arc::new(InMemoryPoisonStore::default()),
|
||||
Some(gate.clone()),
|
||||
None,
|
||||
));
|
||||
let manager = RegionMigrationManager::new(procedure_manager.clone(), env.context_factory());
|
||||
let region_ids = if batch {
|
||||
vec![RegionId::new(1024, 1), RegionId::new(1024, 2)]
|
||||
} else {
|
||||
vec![RegionId::new(1024, 1)]
|
||||
};
|
||||
let region_routes = region_ids
|
||||
.iter()
|
||||
.map(|region_id| RegionRoute {
|
||||
region: Region::new_test(*region_id),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
..Default::default()
|
||||
})
|
||||
.collect();
|
||||
env.create_physical_table_metadata(new_test_table_info(1024), region_routes)
|
||||
.await;
|
||||
|
||||
// A stopped manager must reject submission and release all migration guards.
|
||||
// Starting it then allows the same regions to be submitted successfully.
|
||||
for started in [false, true] {
|
||||
if started {
|
||||
procedure_manager.start().await.unwrap();
|
||||
}
|
||||
let submission = async {
|
||||
if batch {
|
||||
let result = manager
|
||||
.submit_region_migration_task(RegionMigrationTaskBatch {
|
||||
region_ids: region_ids.clone(),
|
||||
from_peer: Peer::empty(1),
|
||||
to_peer: Peer::empty(2),
|
||||
timeout: Duration::from_secs(10),
|
||||
trigger_reason: RegionMigrationTriggerReason::Manual,
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(result.submitted, region_ids);
|
||||
Ok(result.procedure_id)
|
||||
} else {
|
||||
manager
|
||||
.submit_procedure(
|
||||
ProcedureContext::default(),
|
||||
RegionMigrationProcedureTask::new(
|
||||
region_ids[0],
|
||||
Peer::empty(1),
|
||||
Peer::empty(2),
|
||||
Duration::from_secs(10),
|
||||
RegionMigrationTriggerReason::Manual,
|
||||
),
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
tokio::pin!(submission);
|
||||
tokio::select! {
|
||||
biased;
|
||||
result = &mut submission => panic!("Submission returned before registration: {result:?}"),
|
||||
_ = gate.entered.notified() => {}
|
||||
}
|
||||
assert!(futures::poll!(&mut submission).is_pending());
|
||||
assert!(
|
||||
procedure_manager
|
||||
.list_procedures()
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
for region_id in ®ion_ids {
|
||||
assert!(manager.tracker.contains(*region_id));
|
||||
}
|
||||
|
||||
gate.release.notify_one();
|
||||
let result = tokio::time::timeout(Duration::from_secs(5), &mut submission)
|
||||
.await
|
||||
.expect("Submission should return after registration");
|
||||
if started {
|
||||
let procedure_id = result.unwrap().unwrap();
|
||||
assert!(
|
||||
procedure_manager
|
||||
.procedure_state(procedure_id)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
} else {
|
||||
assert_matches!(
|
||||
result.unwrap_err(),
|
||||
error::Error::SubmitProcedure {
|
||||
source: common_procedure::Error::ManagerNotStart { .. },
|
||||
..
|
||||
}
|
||||
);
|
||||
for region_id in ®ion_ids {
|
||||
assert!(!manager.tracker.contains(*region_id));
|
||||
}
|
||||
}
|
||||
}
|
||||
procedure_manager.stop().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_submit_procedure_registration() {
|
||||
check_submission_registration(false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_submit_region_migration_task_registration() {
|
||||
check_submission_registration(true).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_running_procedure() {
|
||||
let env = TestingEnv::new();
|
||||
@@ -1039,6 +1171,7 @@ mod test {
|
||||
#[tokio::test]
|
||||
async fn test_running_regions() {
|
||||
let env = TestingEnv::new();
|
||||
env.procedure_manager().start().await.unwrap();
|
||||
let context_factory = env.context_factory();
|
||||
let manager = RegionMigrationManager::new(env.procedure_manager().clone(), context_factory);
|
||||
let region_id = RegionId::new(1024, 1);
|
||||
@@ -1072,5 +1205,6 @@ mod test {
|
||||
assert_eq!(result.migrating, vec![region_id]);
|
||||
assert_eq!(result.submitted, vec![RegionId::new(1024, 2)]);
|
||||
assert!(result.procedure_id.is_some());
|
||||
env.procedure_manager().stop().await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,6 +101,10 @@ pub(crate) struct RangeResultMemoryLimiter {
|
||||
semaphore: Arc<tokio::sync::Semaphore>,
|
||||
permit_bytes: usize,
|
||||
total_permits: usize,
|
||||
/// Number of acquisitions that found too few permits and parked. Test-only
|
||||
/// signal to synchronize with a caller waiting inside [Self::acquire].
|
||||
#[cfg(test)]
|
||||
waited_acquires: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
|
||||
impl Default for RangeResultMemoryLimiter {
|
||||
@@ -122,6 +126,8 @@ impl RangeResultMemoryLimiter {
|
||||
semaphore: Arc::new(tokio::sync::Semaphore::new(total_permits)),
|
||||
permit_bytes,
|
||||
total_permits,
|
||||
#[cfg(test)]
|
||||
waited_acquires: std::sync::atomic::AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -130,6 +136,12 @@ impl RangeResultMemoryLimiter {
|
||||
self.permit_bytes
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn waited_acquires(&self) -> usize {
|
||||
self.waited_acquires
|
||||
.load(std::sync::atomic::Ordering::Acquire)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn available_permits(&self) -> usize {
|
||||
self.semaphore.available_permits()
|
||||
@@ -146,6 +158,14 @@ impl RangeResultMemoryLimiter {
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
// Nothing awaits between this check and the parking below, so an observed
|
||||
// increment means the caller is about to wait for the missing permits.
|
||||
#[cfg(test)]
|
||||
if self.semaphore.available_permits() < permits {
|
||||
self.waited_acquires
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Release);
|
||||
}
|
||||
|
||||
self.semaphore
|
||||
.acquire_many(permits as u32)
|
||||
.await
|
||||
|
||||
@@ -892,6 +892,7 @@ struct Window {
|
||||
start: Timestamp,
|
||||
end: Timestamp,
|
||||
files: Vec<FileHandle>,
|
||||
/// Start of the half-open window, in seconds.
|
||||
time_window: i64,
|
||||
primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>,
|
||||
}
|
||||
@@ -942,12 +943,7 @@ fn assign_to_windows<'a>(
|
||||
continue;
|
||||
}
|
||||
let (_, end) = f.time_range();
|
||||
let time_window = end
|
||||
.convert_to(TimeUnit::Second)
|
||||
.unwrap()
|
||||
.value()
|
||||
.align_to_ceil_by_bucket(time_window_size)
|
||||
.unwrap_or(i64::MIN);
|
||||
let time_window = time_window_start(end, time_window_size);
|
||||
|
||||
match windows.entry(time_window) {
|
||||
Entry::Occupied(mut e) => {
|
||||
@@ -963,37 +959,41 @@ fn assign_to_windows<'a>(
|
||||
windows.into_iter().collect()
|
||||
}
|
||||
|
||||
/// Uses the same half-open windows as memtables, including timestamps in the first
|
||||
/// second after a boundary. An unrepresentable window start is clamped to `i64::MIN`.
|
||||
fn time_window_start(timestamp: Timestamp, time_window_size: i64) -> i64 {
|
||||
// Conversion to seconds cannot overflow.
|
||||
timestamp
|
||||
.convert_to(TimeUnit::Second)
|
||||
.and_then(|ts| ts.value().align_by_bucket(time_window_size))
|
||||
.unwrap_or(i64::MIN)
|
||||
}
|
||||
|
||||
fn time_window_intersects_range(
|
||||
window_end: i64,
|
||||
window_start: i64,
|
||||
time_window_size: i64,
|
||||
time_range: &TimestampRange,
|
||||
) -> bool {
|
||||
let first_window = match time_range.start() {
|
||||
None => i64::MIN,
|
||||
Some(start) => {
|
||||
let Some(first_window) = start
|
||||
.convert_to(TimeUnit::Second)
|
||||
.and_then(|timestamp| timestamp.value().align_to_ceil_by_bucket(time_window_size))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
first_window
|
||||
}
|
||||
};
|
||||
if time_range.is_empty() {
|
||||
return false;
|
||||
}
|
||||
let first_window = time_range
|
||||
.start()
|
||||
.map_or(i64::MIN, |start| time_window_start(start, time_window_size));
|
||||
let last_window = match time_range.end() {
|
||||
None => i64::MAX,
|
||||
Some(end) => {
|
||||
let Some(last_window) = end
|
||||
// The range end is exclusive; keep the last second that intersects it.
|
||||
let Some(last_second) = end
|
||||
.convert_to_ceil(TimeUnit::Second)
|
||||
.and_then(|timestamp| timestamp.value().checked_sub(1))
|
||||
.and_then(|timestamp| timestamp.align_to_ceil_by_bucket(time_window_size))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
last_window
|
||||
time_window_start(Timestamp::new_second(last_second), time_window_size)
|
||||
}
|
||||
};
|
||||
(first_window..=last_window).contains(&window_end)
|
||||
(first_window..=last_window).contains(&window_start)
|
||||
}
|
||||
|
||||
fn window_has_overlap(this: &Window, windows: &BTreeMap<i64, Window>) -> bool {
|
||||
@@ -1011,7 +1011,7 @@ fn window_has_overlap(this: &Window, windows: &BTreeMap<i64, Window>) -> bool {
|
||||
}
|
||||
|
||||
/// Finds the latest active writing window among all files.
|
||||
/// Returns `None` when there are no files or all files are corrupted.
|
||||
/// Returns `None` when there are no files.
|
||||
fn find_latest_window_in_seconds<'a>(
|
||||
files: impl Iterator<Item = &'a FileHandle>,
|
||||
time_window_size: i64,
|
||||
@@ -1027,9 +1027,7 @@ fn find_latest_window_in_seconds<'a>(
|
||||
latest_timestamp = Some(end);
|
||||
}
|
||||
}
|
||||
latest_timestamp
|
||||
.and_then(|ts| ts.convert_to_ceil(TimeUnit::Second))
|
||||
.and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size))
|
||||
latest_timestamp.map(|ts| time_window_start(ts, time_window_size))
|
||||
}
|
||||
|
||||
/// Finds the active window from the file with the highest sequence number.
|
||||
@@ -1039,9 +1037,7 @@ fn find_latest_window_in_seconds<'a>(
|
||||
/// survives the transient state where level 0 is empty right after its files
|
||||
/// were compacted away. Returns `None` when no file carries a sequence.
|
||||
///
|
||||
/// The window key follows the same convention as [`assign_to_windows`]
|
||||
/// (truncate to seconds, then align up), since it is compared against the
|
||||
/// window keys produced there.
|
||||
/// The window key is its inclusive start, as in [`assign_to_windows`].
|
||||
fn find_active_window_by_sequence<'a>(
|
||||
files: impl Iterator<Item = &'a FileHandle>,
|
||||
time_window_size: i64,
|
||||
@@ -1049,12 +1045,7 @@ fn find_active_window_by_sequence<'a>(
|
||||
files
|
||||
.filter(|f| f.meta_ref().sequence.is_some())
|
||||
.max_by_key(|f| f.meta_ref().sequence)
|
||||
.and_then(|f| {
|
||||
f.time_range()
|
||||
.1
|
||||
.convert_to(TimeUnit::Second)
|
||||
.and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size))
|
||||
})
|
||||
.map(|f| time_window_start(f.time_range().1, time_window_size))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -1347,7 +1338,7 @@ mod tests {
|
||||
new_file_handle_with_sequence(FileId::random(), 0, 999, 1, 0),
|
||||
new_file_handle_with_sequence(FileId::random(), 2000, 2999, 0, 0),
|
||||
];
|
||||
assert_eq!(Some(3), active_window_of(&files));
|
||||
assert_eq!(Some(2), active_window_of(&files));
|
||||
|
||||
// No sequence and no L0: None, exactly as before.
|
||||
let files = [new_file_handle_with_sequence(
|
||||
@@ -1363,7 +1354,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_get_latest_window_in_seconds() {
|
||||
assert_eq!(
|
||||
Some(1),
|
||||
Some(0),
|
||||
find_latest_window_in_seconds([new_file_handle(FileId::random(), 0, 999, 0)].iter(), 1)
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -1375,7 +1366,7 @@ mod tests {
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Some(-9223372036854000),
|
||||
Some(-9223372036857600),
|
||||
find_latest_window_in_seconds(
|
||||
[new_file_handle(FileId::random(), i64::MIN, i64::MIN + 1, 0)].iter(),
|
||||
3600,
|
||||
@@ -1383,7 +1374,7 @@ mod tests {
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
(i64::MAX / 10000000 + 1) * 10000,
|
||||
(i64::MAX / 10000000) * 10000,
|
||||
find_latest_window_in_seconds(
|
||||
[new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0)].iter(),
|
||||
10000,
|
||||
@@ -1392,7 +1383,7 @@ mod tests {
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Some((i64::MAX / 3600000 + 1) * 3600),
|
||||
Some((i64::MAX / 3600000) * 3600),
|
||||
find_latest_window_in_seconds(
|
||||
[
|
||||
new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0),
|
||||
@@ -1404,6 +1395,114 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_window_assignment_and_activity_at_boundaries() {
|
||||
// Before, at, and after positive/negative boundaries, including the first
|
||||
// fractional second: all picker paths must agree with half-open partitions.
|
||||
for (unit, ticks_per_second) in [
|
||||
(TimeUnit::Second, 1),
|
||||
(TimeUnit::Millisecond, 1_000),
|
||||
(TimeUnit::Microsecond, 1_000_000),
|
||||
(TimeUnit::Nanosecond, 1_000_000_000),
|
||||
] {
|
||||
for window_size in [1, 3, 86_400] {
|
||||
for boundary in [-window_size, 0, window_size] {
|
||||
for (offset, expected_window) in [
|
||||
(-1, boundary - window_size),
|
||||
(0, boundary),
|
||||
(ticks_per_second - 1, boundary),
|
||||
(ticks_per_second, boundary + i64::from(window_size == 1)),
|
||||
] {
|
||||
let timestamp = Timestamp::new(boundary * ticks_per_second + offset, unit);
|
||||
let file = FileHandle::new(
|
||||
FileMeta {
|
||||
time_range: (timestamp, timestamp),
|
||||
sequence: NonZeroU64::new(1),
|
||||
..Default::default()
|
||||
},
|
||||
crate::test_util::new_noop_file_purger(),
|
||||
);
|
||||
let files = [file];
|
||||
let windows = assign_to_windows(files.iter(), window_size);
|
||||
assert_eq!(
|
||||
vec![expected_window],
|
||||
windows.keys().copied().collect::<Vec<_>>(),
|
||||
"{timestamp:?}, window size {window_size}"
|
||||
);
|
||||
assert_eq!(
|
||||
Some(expected_window),
|
||||
find_active_window_by_sequence(files.iter(), window_size)
|
||||
);
|
||||
assert_eq!(
|
||||
Some(expected_window),
|
||||
find_latest_window_in_seconds(files.iter(), window_size)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_window_assignment_at_timestamp_limits() {
|
||||
for (seconds, expected_window) in [(i64::MIN, i64::MIN), (i64::MAX, i64::MAX - 1)] {
|
||||
let timestamp = Timestamp::new_second(seconds);
|
||||
let files = [FileHandle::new(
|
||||
FileMeta {
|
||||
time_range: (timestamp, timestamp),
|
||||
sequence: NonZeroU64::new(1),
|
||||
..Default::default()
|
||||
},
|
||||
crate::test_util::new_noop_file_purger(),
|
||||
)];
|
||||
let windows = assign_to_windows(files.iter(), 3);
|
||||
assert_eq!(
|
||||
vec![expected_window],
|
||||
windows.keys().copied().collect::<Vec<_>>()
|
||||
);
|
||||
assert_eq!(
|
||||
Some(expected_window),
|
||||
find_active_window_by_sequence(files.iter(), 3)
|
||||
);
|
||||
assert_eq!(
|
||||
Some(expected_window),
|
||||
find_latest_window_in_seconds(files.iter(), 3)
|
||||
);
|
||||
assert!(time_window_intersects_range(
|
||||
expected_window,
|
||||
3,
|
||||
&TimestampRange::single(timestamp)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_midnight_file_is_not_compacted_with_previous_day() {
|
||||
// Regression for #9441: flush already separates the midnight-only SST.
|
||||
let midnight = 86_400_000;
|
||||
CompactionPickerTestCase {
|
||||
window_size: 86_400,
|
||||
input_files: [
|
||||
(midnight - 240_000, midnight - 180_000),
|
||||
(midnight - 180_000, midnight - 90_000),
|
||||
(midnight - 90_000, midnight - 30_000),
|
||||
(midnight, midnight),
|
||||
]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, (start, end))| {
|
||||
new_file_handle_with_sequence(FileId::random(), start, end, 0, index as u64 + 1)
|
||||
})
|
||||
.collect(),
|
||||
expected_outputs: vec![ExpectedOutput {
|
||||
input_files: vec![0, 1, 2],
|
||||
output_level: 1,
|
||||
}],
|
||||
}
|
||||
.check()
|
||||
.await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_assign_to_windows() {
|
||||
let windows = assign_to_windows(
|
||||
@@ -1433,7 +1532,7 @@ mod tests {
|
||||
assert_eq!(
|
||||
files[0],
|
||||
windows
|
||||
.get(&0)
|
||||
.get(&-3)
|
||||
.unwrap()
|
||||
.files()
|
||||
.next()
|
||||
@@ -1444,7 +1543,7 @@ mod tests {
|
||||
assert_eq!(
|
||||
files[1],
|
||||
windows
|
||||
.get(&3)
|
||||
.get(&0)
|
||||
.unwrap()
|
||||
.files()
|
||||
.next()
|
||||
@@ -1455,7 +1554,7 @@ mod tests {
|
||||
assert_eq!(
|
||||
files[2],
|
||||
windows
|
||||
.get(&12)
|
||||
.get(&9)
|
||||
.unwrap()
|
||||
.files()
|
||||
.next()
|
||||
@@ -1550,6 +1649,7 @@ mod tests {
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let windows = assign_to_windows(file_handles.iter(), time_window);
|
||||
assert_eq!(expected_files.len(), windows.len());
|
||||
|
||||
for (expected_window, overlapping, window_files) in expected_files {
|
||||
let actual_window = windows.get(expected_window).unwrap();
|
||||
@@ -1574,8 +1674,8 @@ mod tests {
|
||||
&[(0, 999), (1000, 1999), (2000, 2999)],
|
||||
2,
|
||||
&[
|
||||
(0, false, vec![(0, 999)]),
|
||||
(2, false, vec![(1000, 1999), (2000, 2999)]),
|
||||
(0, false, vec![(0, 999), (1000, 1999)]),
|
||||
(2, false, vec![(2000, 2999)]),
|
||||
],
|
||||
);
|
||||
|
||||
@@ -1592,9 +1692,8 @@ mod tests {
|
||||
&[(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)],
|
||||
2,
|
||||
&[
|
||||
(0, false, vec![(0, 999)]),
|
||||
(2, false, vec![(1000, 1999), (2000, 2999)]),
|
||||
(4, false, vec![(3000, 3999)]),
|
||||
(0, false, vec![(0, 999), (1000, 1999)]),
|
||||
(2, false, vec![(2000, 2999), (3000, 3999)]),
|
||||
],
|
||||
);
|
||||
|
||||
@@ -1608,9 +1707,8 @@ mod tests {
|
||||
],
|
||||
2,
|
||||
&[
|
||||
(0, true, vec![(0, 999)]),
|
||||
(2, true, vec![(1000, 1999), (2000, 2999)]),
|
||||
(4, true, vec![(0, 3999), (3000, 3999)]),
|
||||
(0, true, vec![(0, 999), (1000, 1999)]),
|
||||
(2, true, vec![(0, 3999), (2000, 2999), (3000, 3999)]),
|
||||
],
|
||||
);
|
||||
|
||||
@@ -1624,43 +1722,39 @@ mod tests {
|
||||
],
|
||||
2,
|
||||
&[
|
||||
(0, false, vec![(0, 999)]),
|
||||
(2, true, vec![(1000, 1999), (2000, 2999)]),
|
||||
(4, true, vec![(1999, 3999), (3000, 3999)]),
|
||||
(0, true, vec![(0, 999), (1000, 1999)]),
|
||||
(2, true, vec![(1999, 3999), (2000, 2999), (3000, 3999)]),
|
||||
],
|
||||
);
|
||||
|
||||
check_assign_to_windows_with_overlapping(
|
||||
&[
|
||||
(0, 999), // window 0
|
||||
(1000, 1999), // window 2
|
||||
(1000, 1999), // window 0
|
||||
(2000, 2999), // window 2
|
||||
(3000, 3999), // window 4
|
||||
(2999, 3999), // window 4
|
||||
(3000, 3999), // window 2
|
||||
(2999, 3999), // window 2
|
||||
],
|
||||
2,
|
||||
&[
|
||||
// window 2 overlaps with window 4
|
||||
(0, false, vec![(0, 999)]),
|
||||
(2, true, vec![(1000, 1999), (2000, 2999)]),
|
||||
(4, true, vec![(2999, 3999), (3000, 3999)]),
|
||||
// Overlap within a window does not imply cross-window overlap.
|
||||
(0, false, vec![(0, 999), (1000, 1999)]),
|
||||
(2, false, vec![(2000, 2999), (2999, 3999), (3000, 3999)]),
|
||||
],
|
||||
);
|
||||
|
||||
check_assign_to_windows_with_overlapping(
|
||||
&[
|
||||
(0, 999), // window 0
|
||||
(1000, 1999), // window 2
|
||||
(1000, 1999), // window 0
|
||||
(2000, 2999), // window 2
|
||||
(3000, 3999), // window 4
|
||||
(0, 1000), // // window 2
|
||||
(3000, 3999), // window 2
|
||||
(0, 2000), // window 2
|
||||
],
|
||||
2,
|
||||
&[
|
||||
// only window 0 overlaps with window 2.
|
||||
(0, true, vec![(0, 999)]),
|
||||
(2, true, vec![(0, 1000), (1000, 1999), (2000, 2999)]),
|
||||
(4, false, vec![(3000, 3999)]),
|
||||
(0, true, vec![(0, 999), (1000, 1999)]),
|
||||
(2, true, vec![(0, 2000), (2000, 2999), (3000, 3999)]),
|
||||
],
|
||||
);
|
||||
}
|
||||
@@ -1680,7 +1774,7 @@ mod tests {
|
||||
new_file_handle_with_size_sequence_and_primary_key_range(
|
||||
FileId::random(),
|
||||
500,
|
||||
1999,
|
||||
2999,
|
||||
0,
|
||||
2,
|
||||
10,
|
||||
@@ -1690,7 +1784,8 @@ mod tests {
|
||||
|
||||
let windows = assign_to_windows(files.iter(), 2);
|
||||
|
||||
let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows);
|
||||
assert_eq!(2, windows.len());
|
||||
let overlapping = window_has_overlap(windows.get(&0).unwrap(), &windows);
|
||||
assert!(!overlapping);
|
||||
}
|
||||
|
||||
@@ -2410,7 +2505,7 @@ mod tests {
|
||||
};
|
||||
|
||||
let output = picker
|
||||
.build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), None)
|
||||
.build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -2427,27 +2522,60 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_filter_time_windows_by_time_range() {
|
||||
let time_range = TimestampRange::new(
|
||||
Timestamp::new_millisecond(1_200),
|
||||
Timestamp::new_millisecond(1_800),
|
||||
)
|
||||
.unwrap();
|
||||
// Interior, exclusive boundary ends, fractional boundary ends, negative
|
||||
// timestamps, and empty ranges must all use the same half-open windows.
|
||||
for (start, end, expected_windows) in [
|
||||
(1_200, 1_800, vec![0]),
|
||||
(0, 3_000, vec![0]),
|
||||
(3_000, 6_000, vec![3]),
|
||||
(2_999, 3_000, vec![0]),
|
||||
(2_999, 3_001, vec![0, 3]),
|
||||
(3_000, 3_001, vec![3]),
|
||||
(3_001, 3_002, vec![3]),
|
||||
(-3_000, 0, vec![-3]),
|
||||
(-3_001, -2_999, vec![-6, -3]),
|
||||
(0, 0, vec![]),
|
||||
(1_500, 1_500, vec![]),
|
||||
] {
|
||||
let time_range = TimestampRange::with_unit(start, end, TimeUnit::Millisecond).unwrap();
|
||||
let selected = [-6, -3, 0, 3, 6]
|
||||
.into_iter()
|
||||
.filter(|window| time_window_intersects_range(*window, 3, &time_range))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(expected_windows, selected, "{time_range}");
|
||||
}
|
||||
}
|
||||
|
||||
assert!(time_window_intersects_range(3, 3, &time_range));
|
||||
assert!(!time_window_intersects_range(9, 3, &time_range));
|
||||
#[test]
|
||||
fn test_filter_time_windows_with_unbounded_and_extreme_ranges() {
|
||||
for (range, expected_windows) in [
|
||||
(TimestampRange::min_to_max(), vec![-3, 0, 3]),
|
||||
(
|
||||
TimestampRange::from_start(Timestamp::new_second(0)),
|
||||
vec![0, 3],
|
||||
),
|
||||
(
|
||||
TimestampRange::until_end(Timestamp::new_second(0), false),
|
||||
vec![-3],
|
||||
),
|
||||
] {
|
||||
let selected = [-3, 0, 3]
|
||||
.into_iter()
|
||||
.filter(|window| time_window_intersects_range(*window, 3, &range))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(expected_windows, selected, "{range}");
|
||||
}
|
||||
|
||||
let boundary_range =
|
||||
TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3)).unwrap();
|
||||
assert!(time_window_intersects_range(0, 3, &boundary_range));
|
||||
assert!(time_window_intersects_range(3, 3, &boundary_range));
|
||||
assert!(!time_window_intersects_range(6, 3, &boundary_range));
|
||||
|
||||
let overflowing_range = TimestampRange::new(
|
||||
let upper_range = TimestampRange::new(
|
||||
Timestamp::new_second(i64::MAX - 1),
|
||||
Timestamp::new_second(i64::MAX),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!time_window_intersects_range(0, 4, &overflowing_range));
|
||||
assert!(!time_window_intersects_range(0, 4, &upper_range));
|
||||
assert!(time_window_intersects_range(i64::MAX - 3, 4, &upper_range));
|
||||
|
||||
let lower_range = TimestampRange::until_end(Timestamp::new_second(i64::MIN), false);
|
||||
assert!(!time_window_intersects_range(i64::MIN, 3, &lower_range));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -2477,7 +2605,7 @@ mod tests {
|
||||
};
|
||||
|
||||
let output = picker
|
||||
.build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), Some(3))
|
||||
.build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), Some(3))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -736,26 +736,33 @@ struct CacheBatchBuffer {
|
||||
buffered_rows: usize,
|
||||
buffered_size: usize,
|
||||
sender: Option<mpsc::UnboundedSender<CacheConcatCommand>>,
|
||||
/// Handle of the background concat task. Cleared once the task owns a finish
|
||||
/// command, so only unfinished scans cancel it.
|
||||
concat_task: Option<common_runtime::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl CacheBatchBuffer {
|
||||
fn new(cache_strategy: &CacheStrategy) -> Self {
|
||||
let sender = cache_strategy.range_result_memory_limiter().map(|limiter| {
|
||||
let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0);
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
common_runtime::spawn_query(run_cache_concat_task(
|
||||
rx,
|
||||
limiter.clone(),
|
||||
skip_threshold_bytes,
|
||||
));
|
||||
tx
|
||||
});
|
||||
let (sender, concat_task) = cache_strategy
|
||||
.range_result_memory_limiter()
|
||||
.map(|limiter| {
|
||||
let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0);
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let task = common_runtime::spawn_query(run_cache_concat_task(
|
||||
rx,
|
||||
limiter.clone(),
|
||||
skip_threshold_bytes,
|
||||
));
|
||||
(tx, task)
|
||||
})
|
||||
.unzip();
|
||||
|
||||
Self {
|
||||
buffered_batches: Vec::new(),
|
||||
buffered_rows: 0,
|
||||
buffered_size: 0,
|
||||
sender,
|
||||
concat_task,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -814,9 +821,22 @@ impl CacheBatchBuffer {
|
||||
part_metrics,
|
||||
result_tx,
|
||||
})
|
||||
.is_err()
|
||||
.is_ok()
|
||||
{
|
||||
self.sender = None;
|
||||
// The task now owns the finish command, so it may keep populating the
|
||||
// cache after the scan stream is dropped.
|
||||
self.concat_task = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CacheBatchBuffer {
|
||||
fn drop(&mut self) {
|
||||
// Still holding the handle means no finish command was sent: the scan was
|
||||
// cancelled or failed and the queued batches can never reach the cache.
|
||||
// Aborting releases them even while the task waits for a memory permit.
|
||||
if let Some(task) = &self.concat_task {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1570,6 +1590,43 @@ mod tests {
|
||||
assert_eq!(replayed[1].num_rows(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() {
|
||||
let strategy = test_cache_strategy();
|
||||
let limiter = strategy.range_result_memory_limiter().unwrap();
|
||||
// Take every permit so the concat task blocks before it can compact.
|
||||
let _permit = limiter
|
||||
.acquire(limiter.available_permits() * limiter.permit_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]);
|
||||
let weak = Arc::downgrade(batch.column(0));
|
||||
let mut buffer = CacheBatchBuffer::new(&strategy);
|
||||
buffer.push(batch.clone()).unwrap();
|
||||
buffer.push(batch).unwrap();
|
||||
// Both batches were handed to the task instead of staying in the buffer.
|
||||
assert!(buffer.buffered_batches.is_empty());
|
||||
|
||||
// Cancel only once the task is parked on the permit, otherwise the abort
|
||||
// could land on a task that was never polled.
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
while limiter.waited_acquires() == 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("concat task must wait for a memory permit");
|
||||
drop(buffer);
|
||||
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
while weak.upgrade().is_some() {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("cancelled cache work must release input without waiting for a permit");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_batch_buffer_finishes_pending_batches() {
|
||||
let strategy = test_cache_strategy();
|
||||
|
||||
@@ -1674,7 +1674,15 @@ impl ScanInput {
|
||||
};
|
||||
match maybe_batch {
|
||||
Some(Ok(batch)) => {
|
||||
let _ = sender.send(Ok(batch)).await;
|
||||
// The receiver is gone when the query is cancelled or finishes early,
|
||||
// so stop reading the source.
|
||||
if let Err(e) = sender.send(Ok(batch)).await {
|
||||
debug!(
|
||||
"Stop parallel scan task, receiver dropped, region_id: {}, error: {}",
|
||||
region_id, e
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
let _ = sender.send(Err(e)).await;
|
||||
@@ -3229,4 +3237,59 @@ mod tests {
|
||||
.build();
|
||||
assert!(exact_sequence_range(&request, &version).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parallel_flat_source_stops_after_receiver_dropped() {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use datatypes::arrow::array::Int64Array;
|
||||
|
||||
const SOURCE_BATCHES: usize = 10_000;
|
||||
|
||||
let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
|
||||
let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
|
||||
let env = SchedulerEnv::new().await;
|
||||
let input = ScanInput::builder(env.access_layer.clone(), mapper).build();
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"v",
|
||||
ArrowDataType::Int64,
|
||||
false,
|
||||
)]));
|
||||
let batch =
|
||||
RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap();
|
||||
let pulled = Arc::new(AtomicUsize::new(0));
|
||||
let mut released = Vec::new();
|
||||
let sources = (0..2)
|
||||
.map(|_| {
|
||||
let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
released.push(release_rx);
|
||||
let batch = batch.clone();
|
||||
let pulled = pulled.clone();
|
||||
futures::stream::iter(0..SOURCE_BATCHES)
|
||||
.map(move |_| {
|
||||
// Dropped together with the source stream.
|
||||
let _ = &release_tx;
|
||||
pulled.fetch_add(1, Ordering::Relaxed);
|
||||
Ok(batch.clone())
|
||||
})
|
||||
.boxed()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let streams = input
|
||||
.create_parallel_flat_sources(sources, Arc::new(Semaphore::new(2)), 1)
|
||||
.unwrap();
|
||||
drop(streams);
|
||||
|
||||
for release_rx in released {
|
||||
tokio::time::timeout(std::time::Duration::from_secs(10), release_rx)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap_err();
|
||||
}
|
||||
// Each task may have read a few batches ahead into its channel before the
|
||||
// receiver went away, but must not drain the rest of its source.
|
||||
assert!(pulled.load(Ordering::Relaxed) < 100);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -833,6 +833,8 @@ impl HistogramFoldStream {
|
||||
self.find_first_complete_bucket(&batch)
|
||||
}
|
||||
|
||||
/// Returns the first group's candidate bucket count once a group boundary is
|
||||
/// observed. Optimistic validation checks its bounds before folding.
|
||||
fn find_first_complete_bucket(&self, batch: &RecordBatch) -> DataFusionResult<Option<usize>> {
|
||||
if batch.num_rows() == 0 {
|
||||
return Ok(None);
|
||||
@@ -840,21 +842,14 @@ impl HistogramFoldStream {
|
||||
|
||||
let vectors = Helper::try_into_vectors(batch.columns())
|
||||
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
|
||||
let le_array = batch.column(self.le_column_index);
|
||||
|
||||
let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len());
|
||||
self.collect_tag_values(&vectors, 0, &mut tag_values_buf);
|
||||
let mut group_start = 0usize;
|
||||
|
||||
for row in 0..batch.num_rows() {
|
||||
for row in 1..batch.num_rows() {
|
||||
if !self.is_same_group(&vectors, row, &tag_values_buf) {
|
||||
// new group begins
|
||||
self.collect_tag_values(&vectors, row, &mut tag_values_buf);
|
||||
group_start = row;
|
||||
}
|
||||
|
||||
if Self::is_positive_infinity(le_array, row) {
|
||||
return Ok(Some(row - group_start + 1));
|
||||
// A new series/timestamp proves that the previous group is
|
||||
// complete, including any numerically equal +Inf boundaries.
|
||||
return Ok(Some(row));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -875,7 +870,10 @@ impl HistogramFoldStream {
|
||||
let field_array = field_array.as_primitive::<Float64Type>();
|
||||
let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len());
|
||||
|
||||
while remaining_rows >= bucket_num && self.mode == FoldMode::Optimistic {
|
||||
// Keep one row of lookahead: +Inf alone does not prove completion.
|
||||
// A group ending at the batch boundary stays buffered until more input
|
||||
// arrives, or flush_remaining completes it at EOF.
|
||||
while remaining_rows > bucket_num && self.mode == FoldMode::Optimistic {
|
||||
self.collect_tag_values(&vectors, cursor, &mut tag_values_buf);
|
||||
if !self.validate_optimistic_group(
|
||||
&vectors,
|
||||
@@ -974,6 +972,9 @@ impl HistogramFoldStream {
|
||||
tag_values: &[ValueRef<'_>],
|
||||
) -> bool {
|
||||
let inf_index = cursor + bucket_num - 1;
|
||||
if self.is_same_group(vectors, inf_index + 1, tag_values) {
|
||||
return false;
|
||||
}
|
||||
if !Self::is_positive_infinity(le_array, inf_index) {
|
||||
return false;
|
||||
}
|
||||
@@ -1261,8 +1262,31 @@ impl HistogramFoldStream {
|
||||
}
|
||||
}
|
||||
|
||||
// check input value
|
||||
if !bucket.windows(2).all(|w| w[0] <= w[1]) {
|
||||
// Detect equal numeric boundaries during validation so unique buckets
|
||||
// need neither another scan nor an allocation.
|
||||
let has_duplicates = !bucket.windows(2).all(|w| w[0] < w[1]);
|
||||
if has_duplicates && !bucket.windows(2).all(|w| w[0] <= w[1]) {
|
||||
return Ok(f64::NAN);
|
||||
}
|
||||
// Counts must be coalesced before repairing monotonicity.
|
||||
let coalesced = has_duplicates.then(|| {
|
||||
let mut bounds = Vec::with_capacity(bucket.len());
|
||||
let mut counts = Vec::with_capacity(counter.len());
|
||||
for (&bound, &count) in bucket.iter().zip(counter) {
|
||||
if bounds.last() == Some(&bound) {
|
||||
*counts.last_mut().unwrap() += count;
|
||||
} else {
|
||||
bounds.push(bound);
|
||||
counts.push(count);
|
||||
}
|
||||
}
|
||||
(bounds, counts)
|
||||
});
|
||||
let (bucket, counter) = match &coalesced {
|
||||
Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()),
|
||||
None => (bucket, counter),
|
||||
};
|
||||
if matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() < 2 {
|
||||
return Ok(f64::NAN);
|
||||
}
|
||||
let counter = match operation {
|
||||
@@ -1326,26 +1350,6 @@ impl HistogramFoldStream {
|
||||
}
|
||||
|
||||
fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 {
|
||||
let coalesced = bucket
|
||||
.windows(2)
|
||||
.any(|bounds| bounds[0] == bounds[1])
|
||||
.then(|| {
|
||||
let mut bounds = Vec::with_capacity(bucket.len());
|
||||
let mut counts = Vec::with_capacity(counter.len());
|
||||
for (&bound, &count) in bucket.iter().zip(counter) {
|
||||
if bounds.last() == Some(&bound) {
|
||||
*counts.last_mut().unwrap() += count;
|
||||
} else {
|
||||
bounds.push(bound);
|
||||
counts.push(count);
|
||||
}
|
||||
}
|
||||
(bounds, counts)
|
||||
});
|
||||
let (bucket, counter) = match &coalesced {
|
||||
Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()),
|
||||
None => (bucket, counter),
|
||||
};
|
||||
let total = *counter.last().unwrap();
|
||||
if total == 0.0 || lower.is_nan() || upper.is_nan() {
|
||||
return f64::NAN;
|
||||
@@ -1453,8 +1457,8 @@ mod test {
|
||||
|
||||
// 12 items
|
||||
let host_column_1 = Arc::new(StringArray::from(vec![
|
||||
"host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1",
|
||||
"host_1", "host_1", "host_1", "host_1",
|
||||
"host_1", "host_1", "host_1", "host_1", "host_1", "host_2", "host_2", "host_2",
|
||||
"host_2", "host_2", "host_3", "host_3",
|
||||
])) as _;
|
||||
let le_column_1 = Arc::new(StringArray::from(vec![
|
||||
"0.001", "0.1", "10", "1000", "+Inf", "0.001", "0.1", "10", "1000", "+inf", "0.001",
|
||||
@@ -1465,14 +1469,14 @@ mod test {
|
||||
])) as _;
|
||||
|
||||
// 2 items
|
||||
let host_column_2 = Arc::new(StringArray::from(vec!["host_1", "host_1"])) as _;
|
||||
let host_column_2 = Arc::new(StringArray::from(vec!["host_3", "host_3"])) as _;
|
||||
let le_column_2 = Arc::new(StringArray::from(vec!["10", "1000"])) as _;
|
||||
let val_column_2 = Arc::new(Float64Array::from(vec![1.0, 1.0])) as _;
|
||||
|
||||
// 11 items
|
||||
let host_column_3 = Arc::new(StringArray::from(vec![
|
||||
"host_1", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2",
|
||||
"host_2", "host_2", "host_2",
|
||||
"host_3", "host_4", "host_4", "host_4", "host_4", "host_4", "host_5", "host_5",
|
||||
"host_5", "host_5", "host_5",
|
||||
])) as _;
|
||||
let le_column_3 = Arc::new(StringArray::from(vec![
|
||||
"+INF", "0.001", "0.1", "10", "1000", "+iNf", "0.001", "0.1", "10", "1000", "+Inf",
|
||||
@@ -1638,10 +1642,10 @@ mod test {
|
||||
| host | val |
|
||||
+--------+-------------------+
|
||||
| host_1 | 257.5 |
|
||||
| host_1 | 5.05 |
|
||||
| host_1 | 0.0004 |
|
||||
| host_2 | NaN |
|
||||
| host_2 | 6.040000000000001 |
|
||||
| host_2 | 5.05 |
|
||||
| host_3 | 0.0004 |
|
||||
| host_4 | NaN |
|
||||
| host_5 | 6.040000000000001 |
|
||||
+--------+-------------------+",
|
||||
);
|
||||
assert_eq!(result_literal, expected);
|
||||
@@ -1818,7 +1822,7 @@ mod test {
|
||||
Field::new("le", DataType::Utf8, true),
|
||||
Field::new("val", DataType::Float64, true),
|
||||
]));
|
||||
let host_column = Arc::new(StringArray::from(vec!["a", "a", "a", "a", "b", "b"])) as _;
|
||||
let host_column = Arc::new(StringArray::from(vec!["a", "a", "b", "b", "c", "c"])) as _;
|
||||
let le_column = Arc::new(StringArray::from(vec![
|
||||
"0.1", "+Inf", "0.1", "1.0", "0.1", "+Inf",
|
||||
])) as _;
|
||||
@@ -1839,8 +1843,8 @@ mod test {
|
||||
| host | val |
|
||||
+------+-----+
|
||||
| a | 0.1 |
|
||||
| a | NaN |
|
||||
| b | 0.1 |
|
||||
| b | NaN |
|
||||
| c | 0.1 |
|
||||
+------+-----+",
|
||||
);
|
||||
assert_eq!(result_literal, expected);
|
||||
@@ -1917,6 +1921,79 @@ mod test {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trailing_invalid_bounds_fall_back_before_eof() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("host", DataType::Utf8, false),
|
||||
Field::new("le", DataType::Utf8, true),
|
||||
Field::new("val", DataType::Float64, false),
|
||||
]));
|
||||
for trailing_bound in [Some("bad"), None] {
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(StringArray::from(vec![
|
||||
"a", "a", "a", "b", "b", "b", "c", "c", "c",
|
||||
])),
|
||||
Arc::new(StringArray::from(
|
||||
[Some("1"), Some("+Inf"), trailing_bound].repeat(3),
|
||||
)),
|
||||
Arc::new(Float64Array::from([2.0, 4.0, 99.0].repeat(3))),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let fold = build_fold_exec_from_batches(vec![batch.clone()], schema.clone(), 0.5, 0);
|
||||
let mut stream = HistogramFoldStream {
|
||||
le_column_index: 1,
|
||||
field_column_index: 2,
|
||||
histogram_column_index: None,
|
||||
operation: fold.operation,
|
||||
normal_indices: vec![0],
|
||||
bucket_size: None,
|
||||
batch_size: 1,
|
||||
output_schema: fold.output_schema.clone(),
|
||||
input_schema: schema.clone(),
|
||||
mode: FoldMode::Optimistic,
|
||||
safe_group: None,
|
||||
input_buffer: vec![],
|
||||
input_buffered_rows: 0,
|
||||
output_buffer: HistogramFoldStream::empty_output_buffer(&schema).unwrap(),
|
||||
output_buffered_rows: 0,
|
||||
input: fold
|
||||
.input
|
||||
.execute(0, SessionContext::default().task_ctx())
|
||||
.unwrap(),
|
||||
metric: BaselineMetrics::new(&fold.metric, 0),
|
||||
};
|
||||
|
||||
// Neither +Inf nor a batch boundary alone completes the group.
|
||||
assert!(stream.fold_input(batch.slice(0, 2)).unwrap().is_none());
|
||||
assert!(stream.fold_input(batch.slice(2, 1)).unwrap().is_none());
|
||||
assert_eq!(stream.mode, FoldMode::Optimistic);
|
||||
|
||||
for (offset, host) in [(3, "a"), (6, "b")] {
|
||||
let output = stream
|
||||
.fold_input(batch.slice(offset, 3))
|
||||
.unwrap()
|
||||
.expect("a completed group should produce output before EOF")
|
||||
.unwrap();
|
||||
assert_eq!(stream.mode, FoldMode::Safe);
|
||||
assert_eq!(stream.input_buffered_rows, 0);
|
||||
assert!(stream.input_buffer.is_empty());
|
||||
assert_eq!(output.num_rows(), 1);
|
||||
assert_eq!(string_array_value_at_index(output.column(0), 0), Some(host));
|
||||
assert_eq!(output.column(1).as_primitive::<Float64Type>().value(0), 1.0);
|
||||
}
|
||||
|
||||
stream.flush_remaining().unwrap();
|
||||
let output = stream.take_output_buf().unwrap().unwrap();
|
||||
assert_eq!(output.num_rows(), 1);
|
||||
assert_eq!(string_array_value_at_index(output.column(0), 0), Some("c"));
|
||||
assert_eq!(output.column(1).as_primitive::<Float64Type>().value(0), 1.0);
|
||||
assert!(stream.take_output_buf().unwrap().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn safe_mode_handles_misaligned_groups() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
@@ -2032,6 +2109,171 @@ mod test {
|
||||
assert!((values[1] - 0.55).abs() < 1e-10, "{values:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluate_quantile_duplicate_bounds() {
|
||||
let cases = [
|
||||
// Issue #9443: coalesce before repairing the count at 1000.
|
||||
(
|
||||
vec![10.0, 10.0, 100.0, 100.0, 1000.0, f64::INFINITY],
|
||||
vec![40.0, 40.0, 50.0, 50.0, 50.0, 100.0],
|
||||
0.95,
|
||||
77.5,
|
||||
),
|
||||
// Repairing the split counts first would incorrectly inflate them.
|
||||
(
|
||||
vec![10.0, 10.0, 100.0, 100.0, f64::INFINITY],
|
||||
vec![30.0, 10.0, 40.0, 20.0, 60.0],
|
||||
0.5,
|
||||
7.5,
|
||||
),
|
||||
(
|
||||
vec![10.0, f64::INFINITY, f64::INFINITY],
|
||||
vec![40.0, 20.0, 30.0],
|
||||
0.5,
|
||||
6.25,
|
||||
),
|
||||
(
|
||||
vec![f64::INFINITY, f64::INFINITY],
|
||||
vec![20.0, 30.0],
|
||||
0.5,
|
||||
f64::NAN,
|
||||
),
|
||||
(
|
||||
vec![-0.0, 0.0, 10.0, f64::INFINITY],
|
||||
vec![10.0, 10.0, 40.0, 40.0],
|
||||
0.75,
|
||||
5.0,
|
||||
),
|
||||
];
|
||||
for (bounds, counts, quantile, expected) in cases {
|
||||
let actual = HistogramFoldStream::evaluate_row(
|
||||
HistogramFoldOperation::Quantile(quantile.into()),
|
||||
&bounds,
|
||||
&counts,
|
||||
)
|
||||
.unwrap();
|
||||
if expected.is_nan() {
|
||||
assert!(actual.is_nan());
|
||||
} else {
|
||||
assert!((actual - expected).abs() < 1e-10, "{actual} != {expected}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_infinity_bounds_across_batches() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new(
|
||||
"ts",
|
||||
DataType::Timestamp(TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("le", DataType::Utf8, false),
|
||||
Field::new("val", DataType::Float64, false),
|
||||
]));
|
||||
let duplicate = [("10", 20.0), ("10.0", 20.0), ("+Inf", 25.0), ("+inf", 25.0)];
|
||||
let regular = [("10", 20.0), ("10.0", 20.0), ("+Inf", 50.0)];
|
||||
// Cover discovery, an already established optimistic bucket count,
|
||||
// fallback, and the final group at EOF.
|
||||
for groups in [
|
||||
vec![duplicate.as_slice()],
|
||||
vec![duplicate.as_slice(), regular.as_slice()],
|
||||
vec![regular.as_slice(), duplicate.as_slice(), regular.as_slice()],
|
||||
] {
|
||||
let mut timestamps = Vec::new();
|
||||
let mut bounds = Vec::new();
|
||||
let mut counts = Vec::new();
|
||||
for (timestamp, group) in groups.iter().enumerate() {
|
||||
for &(bound, count) in *group {
|
||||
timestamps.push(timestamp as i64);
|
||||
bounds.push(bound);
|
||||
counts.push(count);
|
||||
}
|
||||
}
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondArray::from(timestamps)),
|
||||
Arc::new(StringArray::from(bounds)),
|
||||
Arc::new(Float64Array::from(counts)),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
for batch_size in 1..=batch.num_rows() {
|
||||
let batches = (0..batch.num_rows())
|
||||
.step_by(batch_size)
|
||||
.map(|offset| batch.slice(offset, batch_size.min(batch.num_rows() - offset)))
|
||||
.collect();
|
||||
let fold = build_fold_exec_from_batches(batches, schema.clone(), 0.5, 0);
|
||||
let result =
|
||||
datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx())
|
||||
.await
|
||||
.unwrap();
|
||||
let values: Vec<_> = result
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(1)
|
||||
.as_primitive::<Float64Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.copied()
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(values, vec![6.25; groups.len()], "batch size {batch_size}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_bounds_across_batches_and_safe_fallback() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new(
|
||||
"ts",
|
||||
DataType::Timestamp(TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("le", DataType::Utf8, false),
|
||||
Field::new("val", DataType::Float64, false),
|
||||
]));
|
||||
// The first group establishes six buckets for optimistic folding. The
|
||||
// shorter second group forces safe mode; both have duplicate bounds.
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondArray::from(vec![
|
||||
1000, 1000, 1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000, 2000,
|
||||
])),
|
||||
Arc::new(StringArray::from(vec![
|
||||
"10", "10.0", "100", "100.0", "1000", "+Inf", "10", "10.0", "100", "100.0",
|
||||
"+Inf",
|
||||
])),
|
||||
Arc::new(Float64Array::from(vec![
|
||||
40.0, 40.0, 50.0, 50.0, 50.0, 100.0, 40.0, 40.0, 50.0, 50.0, 100.0,
|
||||
])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
// Split each pair of equal boundaries across an input batch boundary.
|
||||
let batches = (0..batch.num_rows()).map(|i| batch.slice(i, 1)).collect();
|
||||
let fold = build_fold_exec_from_batches(batches, schema, 0.95, 0);
|
||||
let result = datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx())
|
||||
.await
|
||||
.unwrap();
|
||||
let values: Vec<_> = result
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(1)
|
||||
.as_primitive::<Float64Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.copied()
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(values, vec![77.5, 77.5]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluate_row_normal_case() {
|
||||
let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY];
|
||||
|
||||
@@ -554,7 +554,9 @@ impl Stream for SeriesDivideStream {
|
||||
}
|
||||
error => return Poll::Ready(error),
|
||||
};
|
||||
self.buffer.push(batch);
|
||||
if batch.num_rows() != 0 {
|
||||
self.buffer.push(batch);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
@@ -624,7 +626,8 @@ impl SeriesDivideStream {
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use datafusion::arrow::array::{
|
||||
DictionaryArray, Int32Array, LargeStringArray, StringArray, StringViewArray, UInt32Array,
|
||||
DictionaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, StringViewArray,
|
||||
UInt32Array,
|
||||
};
|
||||
use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema, UInt32Type};
|
||||
use datafusion::common::ToDFSchema;
|
||||
@@ -1245,4 +1248,216 @@ mod test {
|
||||
assert!(tsid_array.iter().all(|v| v == Some(expected_tsid)));
|
||||
}
|
||||
}
|
||||
|
||||
/// The input stream may legally start with a batch that has zero rows.
|
||||
/// `SeriesDivideStream` must skip it and keep working on the real rows that
|
||||
/// follow (here: two batches of the very same tagged series).
|
||||
#[tokio::test]
|
||||
async fn test_stream_leading_empty_batch_then_same_tags() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("host", DataType::Utf8, false),
|
||||
Field::new("path", DataType::Utf8, false),
|
||||
Field::new(
|
||||
"time_index",
|
||||
DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
]));
|
||||
|
||||
let make_batch = |hosts: Vec<&str>, paths: Vec<&str>, times: Vec<i64>| {
|
||||
RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(StringArray::from(hosts)),
|
||||
Arc::new(StringArray::from(paths)),
|
||||
Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
|
||||
times,
|
||||
)),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
};
|
||||
|
||||
let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
|
||||
MemorySourceConfig::try_new(
|
||||
&[vec![
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
make_batch(
|
||||
vec!["server1", "server1"],
|
||||
vec!["/var/log", "/var/log"],
|
||||
vec![1000, 2000],
|
||||
),
|
||||
make_batch(vec!["server1"], vec!["/var/log"], vec![3000]),
|
||||
]],
|
||||
schema.clone(),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
)));
|
||||
|
||||
let divide_exec = Arc::new(SeriesDivideExec {
|
||||
tag_columns: vec!["host".to_string(), "path".to_string()],
|
||||
time_index_column: "time_index".to_string(),
|
||||
input: memory_exec,
|
||||
metric: ExecutionPlanMetricsSet::new(),
|
||||
});
|
||||
|
||||
let mut stream = divide_exec
|
||||
.execute(0, SessionContext::default().task_ctx())
|
||||
.unwrap();
|
||||
|
||||
let batch = stream.next().await.unwrap().unwrap();
|
||||
assert_eq!(batch.num_rows(), 3);
|
||||
let expected = [
|
||||
("server1", "/var/log", 1000i64),
|
||||
("server1", "/var/log", 2000),
|
||||
("server1", "/var/log", 3000),
|
||||
];
|
||||
for (row, (host, path, time)) in expected.into_iter().enumerate() {
|
||||
assert_eq!(
|
||||
string_array_value_at_index(batch.column(0), row),
|
||||
Some(host)
|
||||
);
|
||||
assert_eq!(
|
||||
string_array_value_at_index(batch.column(1), row),
|
||||
Some(path)
|
||||
);
|
||||
assert_eq!(
|
||||
batch
|
||||
.column(2)
|
||||
.as_any()
|
||||
.downcast_ref::<datafusion::arrow::array::TimestampMillisecondArray>()
|
||||
.unwrap()
|
||||
.value(row),
|
||||
time
|
||||
);
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
/// Several leading empty batches (and one empty batch in the middle of a
|
||||
/// series) must not disturb series division: the two series are still
|
||||
/// emitted separately, including the last one flushed at end of input.
|
||||
#[tokio::test]
|
||||
async fn test_stream_repeated_empty_batches_and_two_series() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("host", DataType::Utf8, false),
|
||||
Field::new("value", DataType::Int64, false),
|
||||
Field::new(
|
||||
"time_index",
|
||||
DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
]));
|
||||
|
||||
let make_batch = |hosts: Vec<&str>, values: Vec<i64>| {
|
||||
RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(StringArray::from(hosts)),
|
||||
Arc::new(Int64Array::from(values.clone())),
|
||||
Arc::new(datafusion::arrow::array::TimestampMillisecondArray::from(
|
||||
values.iter().map(|v| v * 1000).collect::<Vec<_>>(),
|
||||
)),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
};
|
||||
|
||||
let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
|
||||
MemorySourceConfig::try_new(
|
||||
&[vec![
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
make_batch(vec!["server1", "server1"], vec![1, 2]),
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
make_batch(vec!["server1"], vec![3]),
|
||||
make_batch(vec!["server2", "server2"], vec![4, 5]),
|
||||
]],
|
||||
schema.clone(),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
)));
|
||||
|
||||
let divide_exec = Arc::new(SeriesDivideExec {
|
||||
tag_columns: vec!["host".to_string()],
|
||||
time_index_column: "time_index".to_string(),
|
||||
input: memory_exec,
|
||||
metric: ExecutionPlanMetricsSet::new(),
|
||||
});
|
||||
|
||||
let mut stream = divide_exec
|
||||
.execute(0, SessionContext::default().task_ctx())
|
||||
.unwrap();
|
||||
|
||||
for (expected_host, expected_values) in
|
||||
[("server1", vec![1i64, 2, 3]), ("server2", vec![4, 5])]
|
||||
{
|
||||
let batch = stream.next().await.unwrap().unwrap();
|
||||
assert_eq!(batch.num_rows(), expected_values.len());
|
||||
for (row, value) in expected_values.into_iter().enumerate() {
|
||||
assert_eq!(
|
||||
string_array_value_at_index(batch.column(0), row),
|
||||
Some(expected_host)
|
||||
);
|
||||
assert_eq!(
|
||||
batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap()
|
||||
.value(row),
|
||||
value
|
||||
);
|
||||
}
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
/// A stream that only ever yields zero-row batches must terminate with no
|
||||
/// batches instead of failing.
|
||||
#[tokio::test]
|
||||
async fn test_stream_all_empty_batches_yields_nothing() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("host", DataType::Utf8, false),
|
||||
Field::new("value", DataType::Int64, false),
|
||||
Field::new(
|
||||
"time_index",
|
||||
DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
]));
|
||||
|
||||
let memory_exec: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
|
||||
MemorySourceConfig::try_new(
|
||||
&[vec![
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
RecordBatch::new_empty(schema.clone()),
|
||||
]],
|
||||
schema.clone(),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
)));
|
||||
|
||||
let divide_exec = Arc::new(SeriesDivideExec {
|
||||
tag_columns: vec!["host".to_string()],
|
||||
time_index_column: "time_index".to_string(),
|
||||
input: memory_exec,
|
||||
metric: ExecutionPlanMetricsSet::new(),
|
||||
});
|
||||
|
||||
let mut stream = divide_exec
|
||||
.execute(0, SessionContext::default().task_ctx())
|
||||
.unwrap();
|
||||
assert!(stream.next().await.is_none());
|
||||
|
||||
let collected =
|
||||
datafusion::physical_plan::collect(divide_exec, SessionContext::default().task_ctx())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(collected.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -205,10 +205,14 @@ impl PhysicalOptimizerRule for MatchesConstantTermOptimizer {
|
||||
})?;
|
||||
|
||||
if new_pred.transformed {
|
||||
let exec = FilterExecBuilder::new(new_pred.data, filter.input().clone())
|
||||
.with_default_selectivity(filter.default_selectivity())
|
||||
.apply_projection_by_ref(filter.projection().as_ref())
|
||||
.and_then(|x| x.build())?;
|
||||
// `FilterExecBuilder::from(&FilterExec)` already carries over the
|
||||
// input, projection, default selectivity, batch size and fetch, so
|
||||
// only the predicate has to be replaced. Re-applying the projection
|
||||
// is not only pointless but wrong: `apply_projection` composes with
|
||||
// the projection copied from the builder.
|
||||
let exec = FilterExecBuilder::from(filter)
|
||||
.with_predicate(new_pred.data)
|
||||
.build()?;
|
||||
return Ok(Transformed::yes(Arc::new(exec) as _));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::any::Any;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
@@ -202,6 +203,13 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Time out of the representable millisecond range: {:?}", time))]
|
||||
SystemTimeOutOfRange {
|
||||
time: SystemTime,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("vector cannot contain metrics with the same labelset"))]
|
||||
SameLabelSet {
|
||||
#[snafu(implicit)]
|
||||
@@ -243,6 +251,7 @@ impl ErrorExt for Error {
|
||||
| UnsupportedMatcherOp { .. }
|
||||
| SameLabelSet { .. }
|
||||
| TimestampOutOfRange { .. }
|
||||
| SystemTimeOutOfRange { .. }
|
||||
| InvalidRegularExpression { .. }
|
||||
| InvalidDestinationLabelName { .. } => StatusCode::InvalidArguments,
|
||||
|
||||
|
||||
@@ -23,9 +23,31 @@ use snafu::{OptionExt, ResultExt};
|
||||
use table::TableRef;
|
||||
|
||||
use crate::promql::error::{
|
||||
DataFusionPlanningSnafu, Result, TimeIndexNotFoundSnafu, TimestampOutOfRangeSnafu,
|
||||
DataFusionPlanningSnafu, Result, SystemTimeOutOfRangeSnafu, TimeIndexNotFoundSnafu,
|
||||
TimestampOutOfRangeSnafu,
|
||||
};
|
||||
|
||||
/// Converts a [`SystemTime`] to a millisecond [`Timestamp`].
|
||||
///
|
||||
/// `duration_since` reports an instant before the epoch as an error rather than
|
||||
/// a negative duration, and an RFC3339 query parameter can name one, so the sign
|
||||
/// is recovered here instead of unwrapping.
|
||||
fn millis_since_epoch(time: SystemTime) -> Result<Timestamp> {
|
||||
let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) {
|
||||
Ok(duration) => (duration.as_millis(), false),
|
||||
Err(earlier) => (earlier.duration().as_millis(), true),
|
||||
};
|
||||
let millis = i64::try_from(millis)
|
||||
.ok()
|
||||
.with_context(|| SystemTimeOutOfRangeSnafu { time })?;
|
||||
|
||||
Ok(Timestamp::new_millisecond(if before_epoch {
|
||||
-millis
|
||||
} else {
|
||||
millis
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr {
|
||||
time_index_expr
|
||||
.clone()
|
||||
@@ -67,14 +89,12 @@ pub fn rewrite_label_values_query(
|
||||
})?;
|
||||
|
||||
// We only support millisecond precision at most.
|
||||
let start =
|
||||
Timestamp::new_millisecond(start.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64);
|
||||
let start = millis_since_epoch(start)?;
|
||||
let start = start.convert_to(unit).context(TimestampOutOfRangeSnafu {
|
||||
timestamp: start.value(),
|
||||
unit,
|
||||
})?;
|
||||
let end =
|
||||
Timestamp::new_millisecond(end.duration_since(UNIX_EPOCH).unwrap().as_millis() as i64);
|
||||
let end = millis_since_epoch(end)?;
|
||||
let end = end.convert_to(unit).context(TimestampOutOfRangeSnafu {
|
||||
timestamp: end.value(),
|
||||
unit,
|
||||
@@ -98,3 +118,28 @@ pub fn rewrite_label_values_query(
|
||||
|
||||
Ok(logical_plan)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn millis_before_the_epoch_are_negative() {
|
||||
// `SystemTime::duration_since` reports these as an error; unwrapping it
|
||||
// panicked on any request whose RFC3339 start named a pre-epoch instant.
|
||||
let time = UNIX_EPOCH - Duration::from_millis(1);
|
||||
assert_eq!(millis_since_epoch(time).unwrap().value(), -1);
|
||||
assert_eq!(millis_since_epoch(UNIX_EPOCH).unwrap().value(), 0);
|
||||
|
||||
let time = UNIX_EPOCH + Duration::from_millis(1);
|
||||
assert_eq!(millis_since_epoch(time).unwrap().value(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn millis_beyond_i64_are_rejected() {
|
||||
let time = UNIX_EPOCH + Duration::from_secs(1 << 60);
|
||||
assert!(millis_since_epoch(time).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1554,6 +1554,18 @@ impl PromPlanner {
|
||||
}
|
||||
} else if right_is_empty_metric && !left_is_empty_metric {
|
||||
self.ctx = left_context.clone();
|
||||
} else if rhs.value_type() == ValueType::Scalar
|
||||
&& lhs.value_type() == ValueType::Vector
|
||||
{
|
||||
// A computed scalar on the right (`time()`, `scalar(...)`) has no labels;
|
||||
// the result keeps the labels of the vector on the left.
|
||||
self.ctx = left_context.clone();
|
||||
}
|
||||
// An operand without a table name, such as a comparison filter's projection, is
|
||||
// joined under the `""` qualifier; resolve the projected columns against it.
|
||||
// Otherwise a time index present on both sides is ambiguous.
|
||||
if self.ctx.table_name.is_none() {
|
||||
self.ctx.table_name = Some(String::new());
|
||||
}
|
||||
// Computed scalars reach this join path instead of the literal projection paths.
|
||||
// Broadcast them for arithmetic in the same way as literal scalars.
|
||||
@@ -2596,13 +2608,17 @@ impl PromPlanner {
|
||||
|
||||
let accepts_empty = matcher.is_match("");
|
||||
let column_name = Self::find_case_sensitive_column(table_schema, matcher.name.as_str());
|
||||
// Prometheus reads a label a series does not carry as the empty
|
||||
// string. A row can miss a label two ways: the table has no column
|
||||
// for it, or the column exists but is NULL on that row — the latter
|
||||
// is the norm for logical metrics sharing a physical table, which
|
||||
// holds the union of their label columns.
|
||||
let col = if let Some(column_name) = column_name {
|
||||
let column = DfExpr::Column(Column::from_name(&column_name));
|
||||
let field = table_schema
|
||||
.index_of_column_by_name(None, &column_name)
|
||||
.map(|index| table_schema.field(index));
|
||||
if accepts_empty
|
||||
&& column_name == OTLP_AGGREGATION_TEMPORALITY_LABEL
|
||||
&& let Some(data_type) = field
|
||||
.filter(|field| {
|
||||
field.is_nullable()
|
||||
@@ -2621,7 +2637,6 @@ impl PromPlanner {
|
||||
}
|
||||
} else {
|
||||
DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None)
|
||||
.alias(matcher.name.clone())
|
||||
};
|
||||
let lit = DfExpr::Literal(ScalarValue::Utf8(Some(matcher.value)), None);
|
||||
let expr = match matcher.op {
|
||||
|
||||
@@ -521,7 +521,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn temporality_matchers_treat_null_as_absent() {
|
||||
async fn matchers_read_absent_labels_as_empty() {
|
||||
let marker = OTLP_AGGREGATION_TEMPORALITY_LABEL;
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
marker,
|
||||
@@ -593,31 +593,82 @@ async fn temporality_matchers_treat_null_as_absent() {
|
||||
);
|
||||
}
|
||||
|
||||
let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
// The rule is about NULL, not about the marker: any nullable label column
|
||||
// can be NULL where the series does not carry the label. A non-nullable one
|
||||
// has nothing to normalize.
|
||||
for (nullable, wants_coalesce) in [(true, true), (false, false)] {
|
||||
let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"label",
|
||||
ArrowDataType::Utf8,
|
||||
nullable,
|
||||
)]));
|
||||
let ordinary_scan = LogicalPlanBuilder::scan(
|
||||
"ordinary_labels",
|
||||
provider_as_source(Arc::new(
|
||||
MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(),
|
||||
)),
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
let PromExpr::VectorSelector(selector) =
|
||||
parser::parse(r#"metric{label!="delta"}"#).unwrap()
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema())
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.join(" AND ");
|
||||
assert_eq!(
|
||||
wants_coalesce,
|
||||
expressions.contains("coalesce"),
|
||||
"nullable={nullable}: {expressions}"
|
||||
);
|
||||
}
|
||||
|
||||
// The other way a label goes absent is having no column for it at all. The
|
||||
// literal standing in for the label has to be usable as a predicate.
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"label",
|
||||
ArrowDataType::Utf8,
|
||||
true,
|
||||
)]));
|
||||
let ordinary_scan = LogicalPlanBuilder::scan(
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(StringArray::from(vec![Some("value")]))],
|
||||
)
|
||||
.unwrap();
|
||||
let scan = LogicalPlanBuilder::scan(
|
||||
"ordinary_labels",
|
||||
provider_as_source(Arc::new(
|
||||
MemTable::try_new(ordinary_schema, vec![vec![]]).unwrap(),
|
||||
MemTable::try_new(schema, vec![vec![batch]]).unwrap(),
|
||||
)),
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{label!="delta"}"#).unwrap()
|
||||
let PromExpr::VectorSelector(selector) = parser::parse(r#"metric{absent!="delta"}"#).unwrap()
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
let expressions = PromPlanner::matchers_to_expr(selector.matchers, ordinary_scan.schema())
|
||||
let expressions = PromPlanner::matchers_to_expr(selector.matchers, scan.schema()).unwrap();
|
||||
assert_eq!(
|
||||
r#"Utf8("") != Utf8("delta")"#,
|
||||
expressions.iter().map(ToString::to_string).join(" AND ")
|
||||
);
|
||||
|
||||
let plan = LogicalPlanBuilder::from(scan)
|
||||
.filter(conjunction(expressions).unwrap())
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.join(" AND ");
|
||||
assert!(!expressions.contains("coalesce"), "{expressions}");
|
||||
.build()
|
||||
.unwrap();
|
||||
let (_, batches) = execute(plan, &build_query_engine_state()).await;
|
||||
let rows = batches.iter().map(|batch| batch.num_rows()).sum::<usize>();
|
||||
assert_eq!(1, rows);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -25,7 +25,7 @@ pub fn resolve_addr(bind_addr: &str, server_addr: Option<&str>) -> String {
|
||||
} else {
|
||||
// otherwise, resolve port from bind_addr
|
||||
// should be safe to unwrap here because bind_addr is already validated
|
||||
let port = bind_addr.split(':').nth(1).unwrap();
|
||||
let port = bind_addr.rsplit_once(':').unwrap().1;
|
||||
format!("{server_addr}:{port}")
|
||||
}
|
||||
}
|
||||
@@ -50,6 +50,16 @@ mod tests {
|
||||
super::resolve_addr("127.0.0.1:3001", Some("tomcat:3002"))
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
"tomcat:3001",
|
||||
super::resolve_addr("[::]:3001", Some("tomcat"))
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
"tomcat:3001",
|
||||
super::resolve_addr("[2001:db8::10]:3001", Some("tomcat"))
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
"127.0.0.1:3001",
|
||||
super::resolve_addr("127.0.0.1:3001", None)
|
||||
|
||||
@@ -49,13 +49,13 @@ use futures::future::join_all;
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use itertools::Itertools;
|
||||
use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers};
|
||||
use promql_parser::parser::token::{self};
|
||||
use promql_parser::parser::token::TokenType;
|
||||
use promql_parser::parser::value::ValueType;
|
||||
use promql_parser::parser::{
|
||||
AggregateExpr, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, MatrixSelector, ParenExpr,
|
||||
SubqueryExpr, UnaryExpr, VectorSelector,
|
||||
AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier,
|
||||
MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector,
|
||||
};
|
||||
use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryStatement};
|
||||
use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement};
|
||||
use serde::de::{self, MapAccess, Visitor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
@@ -74,7 +74,7 @@ use table::requests::{
|
||||
pub use super::result::prometheus_resp::{PromSampleValue, PrometheusJsonResponse};
|
||||
use crate::error::{
|
||||
CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu,
|
||||
NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu,
|
||||
NotSupportedSnafu, ParseTimestampSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu,
|
||||
};
|
||||
use crate::http::header::collect_plan_metrics;
|
||||
use crate::otlp::metrics::ucum_to_openmetrics_unit;
|
||||
@@ -1373,6 +1373,72 @@ fn promql_expr_to_metric_name(expr: &PromqlExpr) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the metric name survives a binary operation; see
|
||||
/// [`binary_metric_name_side`] for the operand that supplies it.
|
||||
///
|
||||
/// Follows Prometheus' `shouldDropMetricName` and `resultMetric`: set operators
|
||||
/// and comparisons keep the name, while arithmetic operators drop it. A
|
||||
/// comparison also drops it when it returns a bool or is a one-to-one
|
||||
/// `on(...)` match that does not keep `__name__` or an `ignoring(...)` match
|
||||
/// that excludes it. Prometheus keeps only the `on(...)` labels or deletes the
|
||||
/// `ignoring(...)` labels for one-to-one matches.
|
||||
fn binary_keeps_metric_name(op: &TokenType, modifier: Option<&BinModifier>) -> bool {
|
||||
if op.is_set_operator() {
|
||||
return true;
|
||||
}
|
||||
|
||||
if !op.is_comparison_operator() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let Some(m) = modifier else {
|
||||
return true;
|
||||
};
|
||||
|
||||
if m.return_bool {
|
||||
return false;
|
||||
}
|
||||
|
||||
if m.card == VectorMatchCardinality::OneToOne {
|
||||
match &m.matching {
|
||||
Some(LabelModifier::Include(labels))
|
||||
if !labels.labels.contains(&METRIC_NAME.to_string()) =>
|
||||
{
|
||||
return false;
|
||||
}
|
||||
Some(LabelModifier::Exclude(labels))
|
||||
if labels.labels.contains(&METRIC_NAME.to_string()) =>
|
||||
{
|
||||
return false;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// The operand whose labels supply the result metric name of a binary
|
||||
/// operation.
|
||||
///
|
||||
/// Prometheus builds the result from the side whose labels it carries, i.e. the
|
||||
/// "many" side of the match: `group_right` makes the right-hand operand that
|
||||
/// side, and a scalar left-hand operand (for example `0.5 < a`) leaves the
|
||||
/// vector on the right. Every other case takes the left-hand operand.
|
||||
fn binary_metric_name_side<'a>(
|
||||
lhs: &'a PromqlExpr,
|
||||
rhs: &'a PromqlExpr,
|
||||
modifier: Option<&BinModifier>,
|
||||
) -> &'a PromqlExpr {
|
||||
let group_right =
|
||||
modifier.is_some_and(|m| matches!(m.card, VectorMatchCardinality::OneToMany(_)));
|
||||
if group_right || matches!(lhs.value_type(), ValueType::Scalar) {
|
||||
rhs
|
||||
} else {
|
||||
lhs
|
||||
}
|
||||
}
|
||||
|
||||
/// Recursively collect all metric names from a PromQL expression
|
||||
fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet<String>) {
|
||||
match expr {
|
||||
@@ -1395,14 +1461,16 @@ fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet<String>) {
|
||||
collect_metric_names(expr, metric_names)
|
||||
}
|
||||
PromqlExpr::Unary(UnaryExpr { .. }) => metric_names.clear(),
|
||||
PromqlExpr::Binary(BinaryExpr { lhs, op, .. }) => {
|
||||
if matches!(
|
||||
op.id(),
|
||||
token::T_LAND // INTERSECT
|
||||
| token::T_LOR // UNION
|
||||
| token::T_LUNLESS // EXCEPT
|
||||
) {
|
||||
collect_metric_names(lhs, metric_names)
|
||||
PromqlExpr::Binary(BinaryExpr {
|
||||
lhs,
|
||||
op,
|
||||
rhs,
|
||||
modifier,
|
||||
..
|
||||
}) => {
|
||||
if binary_keeps_metric_name(op, modifier.as_ref()) {
|
||||
let side = binary_metric_name_side(lhs, rhs, modifier.as_ref());
|
||||
collect_metric_names(side, metric_names)
|
||||
} else {
|
||||
metric_names.clear()
|
||||
}
|
||||
@@ -1711,9 +1779,35 @@ pub async fn label_values_query(
|
||||
);
|
||||
let catalog_manager = handler.catalog_manager();
|
||||
|
||||
let mut table_names = try_call_return_response!(
|
||||
retrieve_table_names(&query_ctx, catalog_manager, matches).await
|
||||
);
|
||||
// An empty `match[]` enumerates every metric; otherwise only the
|
||||
// selectors answerable from metadata go down that path.
|
||||
let enumerate_all = matches.is_empty();
|
||||
let (metadata_selectors, label_selectors) =
|
||||
try_call_return_response!(split_selectors_by_label_use(&matches));
|
||||
|
||||
let mut table_names = if enumerate_all || !metadata_selectors.is_empty() {
|
||||
try_call_return_response!(
|
||||
retrieve_table_names(&query_ctx, catalog_manager, metadata_selectors).await
|
||||
)
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
if !label_selectors.is_empty() {
|
||||
table_names.extend(try_call_return_response!(
|
||||
retrieve_table_names_by_labels(
|
||||
&handler,
|
||||
label_selectors,
|
||||
params.start.as_deref(),
|
||||
params.end.as_deref(),
|
||||
&query_ctx,
|
||||
)
|
||||
.await
|
||||
));
|
||||
table_names.sort_unstable();
|
||||
table_names.dedup();
|
||||
}
|
||||
|
||||
table_names = try_call_return_response!(
|
||||
handler
|
||||
.filter_metadata_metric_names(
|
||||
@@ -1897,6 +1991,131 @@ fn take_metric_name(selector: &mut VectorSelector) -> Option<String> {
|
||||
Some(name)
|
||||
}
|
||||
|
||||
/// Removes every `__name__` matcher from the selector and returns them, so the
|
||||
/// rest can be planned as column predicates. A name given as `VectorSelector::name`
|
||||
/// comes back as an equality matcher, making both spellings filter alike.
|
||||
fn take_metric_name_matchers(selector: &mut VectorSelector) -> Vec<Matcher> {
|
||||
let mut taken = Vec::new();
|
||||
if let Some(name) = selector.name.take() {
|
||||
taken.push(Matcher::new(MatchOp::Equal, METRIC_NAME_LABEL, &name));
|
||||
}
|
||||
|
||||
let (name_matchers, rest) = std::mem::take(&mut selector.matchers.matchers)
|
||||
.into_iter()
|
||||
.partition(|matcher| matcher.name == METRIC_NAME_LABEL);
|
||||
selector.matchers.matchers = rest;
|
||||
taken.extend(name_matchers);
|
||||
|
||||
taken
|
||||
}
|
||||
|
||||
/// Whether a metric name satisfies every `__name__` matcher of one selector.
|
||||
///
|
||||
/// Negated matchers are honoured here, unlike in [`retrieve_table_names`] where
|
||||
/// they keep every table so the caller authorizes the full candidate set: the
|
||||
/// names reaching this point are already narrowed by the data, so filtering them
|
||||
/// can only remove names, never widen what a caller gets to see.
|
||||
fn metric_name_matches(table_name: &str, matchers: &[Matcher]) -> bool {
|
||||
matchers.iter().all(|matcher| match &matcher.op {
|
||||
MatchOp::Equal => table_name == matcher.value,
|
||||
MatchOp::NotEqual => table_name != matcher.value,
|
||||
MatchOp::Re(re) => re.is_match(table_name),
|
||||
MatchOp::NotRe(re) => !re.is_match(table_name),
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether a matcher constrains an ordinary label. The others name the metric,
|
||||
/// the database or the field, none of which is a column to scan.
|
||||
fn is_ordinary_label_matcher(matcher: &Matcher) -> bool {
|
||||
matcher.name != METRIC_NAME_LABEL
|
||||
&& matcher.name != FIELD_NAME_LABEL
|
||||
&& !is_database_selection_label(&matcher.name)
|
||||
}
|
||||
|
||||
/// Splits `match[]` selectors by whether they constrain an ordinary label. The
|
||||
/// first group is answerable from table metadata; the second needs the data read
|
||||
/// and is returned as parsed selectors.
|
||||
///
|
||||
/// `or` matchers stay in the metadata group, which ignores them, rather than
|
||||
/// being silently dropped from a data scan that cannot express them.
|
||||
fn split_selectors_by_label_use(matches: &[String]) -> Result<(Vec<String>, Vec<VectorSelector>)> {
|
||||
let mut metadata_only = Vec::new();
|
||||
let mut with_labels = Vec::new();
|
||||
|
||||
for selector in matches {
|
||||
let expr = promql_parser::parser::parse(selector)
|
||||
.map_err(|reason| InvalidQuerySnafu { reason }.build())?;
|
||||
let PromqlExpr::VectorSelector(vector_selector) = expr else {
|
||||
return InvalidQuerySnafu {
|
||||
reason: "expected vector selector".to_string(),
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
|
||||
let constrains_labels = vector_selector.matchers.or_matchers.is_empty()
|
||||
&& vector_selector
|
||||
.matchers
|
||||
.matchers
|
||||
.iter()
|
||||
.any(is_ordinary_label_matcher);
|
||||
if constrains_labels {
|
||||
with_labels.push(vector_selector);
|
||||
} else {
|
||||
metadata_only.push(selector.clone());
|
||||
}
|
||||
}
|
||||
|
||||
Ok((metadata_only, with_labels))
|
||||
}
|
||||
|
||||
/// Resolves selectors constraining ordinary labels into metric names: the data
|
||||
/// answers the label matchers, then each selector's `__name__` matchers narrow
|
||||
/// the names it found.
|
||||
async fn retrieve_table_names_by_labels(
|
||||
handler: &PrometheusHandlerRef,
|
||||
selectors: Vec<VectorSelector>,
|
||||
start: Option<&str>,
|
||||
end: Option<&str>,
|
||||
query_ctx: &QueryContextRef,
|
||||
) -> Result<Vec<String>> {
|
||||
let start_arg = start.map(str::to_string).unwrap_or_else(yesterday_rfc3339);
|
||||
let end_arg = end.map(str::to_string).unwrap_or_else(current_time_rfc3339);
|
||||
let start = QueryLanguageParser::parse_promql_timestamp(&start_arg).with_context(|_| {
|
||||
ParseTimestampSnafu {
|
||||
timestamp: start_arg.clone(),
|
||||
}
|
||||
})?;
|
||||
let end = QueryLanguageParser::parse_promql_timestamp(&end_arg).with_context(|_| {
|
||||
ParseTimestampSnafu {
|
||||
timestamp: end_arg.clone(),
|
||||
}
|
||||
})?;
|
||||
|
||||
let schema = query_ctx.current_schema();
|
||||
let mut table_names = Vec::new();
|
||||
for mut selector in selectors {
|
||||
let name_matchers = take_metric_name_matchers(&mut selector);
|
||||
// The database and field matchers name no column, and the metadata path
|
||||
// ignores them too.
|
||||
let label_matchers = selector
|
||||
.matchers
|
||||
.matchers
|
||||
.into_iter()
|
||||
.filter(is_ordinary_label_matcher)
|
||||
.collect();
|
||||
let matched = handler
|
||||
.query_metric_names_by_labels(label_matchers, &schema, start, end, query_ctx)
|
||||
.await?;
|
||||
table_names.extend(
|
||||
matched
|
||||
.into_iter()
|
||||
.filter(|name| metric_name_matches(name, &name_matchers)),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(table_names)
|
||||
}
|
||||
|
||||
async fn retrieve_table_names(
|
||||
query_ctx: &QueryContext,
|
||||
catalog_manager: CatalogManagerRef,
|
||||
@@ -2407,6 +2626,10 @@ mod tests {
|
||||
deny_operation: bool,
|
||||
denied_table: Option<&'static str>,
|
||||
metric_names: Vec<String>,
|
||||
/// Names the label-matcher path resolves, kept apart from `metric_names`
|
||||
/// so a test can tell which path answered.
|
||||
label_metric_names: Vec<String>,
|
||||
label_lookups: Mutex<Vec<Vec<Matcher>>>,
|
||||
queries: Mutex<Vec<String>>,
|
||||
ordered_outputs: Mutex<Vec<bool>>,
|
||||
}
|
||||
@@ -2482,6 +2705,18 @@ mod tests {
|
||||
Ok(self.metric_names.clone())
|
||||
}
|
||||
|
||||
async fn query_metric_names_by_labels(
|
||||
&self,
|
||||
matchers: Vec<Matcher>,
|
||||
_: &str,
|
||||
_: std::time::SystemTime,
|
||||
_: std::time::SystemTime,
|
||||
_: &QueryContextRef,
|
||||
) -> Result<Vec<String>> {
|
||||
self.label_lookups.lock().unwrap().push(matchers);
|
||||
Ok(self.label_metric_names.clone())
|
||||
}
|
||||
|
||||
async fn query_label_values(
|
||||
&self,
|
||||
_: String,
|
||||
@@ -2547,6 +2782,8 @@ mod tests {
|
||||
deny_operation: false,
|
||||
denied_table: None,
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
});
|
||||
@@ -2600,6 +2837,8 @@ mod tests {
|
||||
deny_operation: false,
|
||||
denied_table: None,
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
});
|
||||
@@ -2684,6 +2923,8 @@ mod tests {
|
||||
deny_operation: false,
|
||||
denied_table: Some("denied"),
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
})),
|
||||
@@ -2700,6 +2941,153 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// A handler over the logical metric tables `cpu_user` and `cpu_system`,
|
||||
/// which is what the metadata path enumerates, with the label-matcher path
|
||||
/// answering `label_metric_names`.
|
||||
fn label_values_handler(label_metric_names: Vec<&str>) -> Arc<TestPrometheusHandler> {
|
||||
let mut cpu_user = test_table_info(
|
||||
1024,
|
||||
"cpu_user",
|
||||
DEFAULT_SCHEMA_NAME,
|
||||
DEFAULT_CATALOG_NAME,
|
||||
Arc::new(Schema::new(vec![])),
|
||||
);
|
||||
cpu_user.meta.options.extra_options.insert(
|
||||
LOGICAL_TABLE_METADATA_KEY.to_string(),
|
||||
"physical_metrics".to_string(),
|
||||
);
|
||||
let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&cpu_user));
|
||||
let mut cpu_system = cpu_user.clone();
|
||||
cpu_system.ident.table_id = 1025;
|
||||
cpu_system.name = "cpu_system".to_string();
|
||||
manager
|
||||
.register_table_sync(RegisterTableRequest {
|
||||
catalog: DEFAULT_CATALOG_NAME.to_string(),
|
||||
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
||||
table_name: cpu_system.name.clone(),
|
||||
table_id: cpu_system.table_id(),
|
||||
table: EmptyTable::from_table_info(&cpu_system),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
Arc::new(TestPrometheusHandler {
|
||||
catalog_manager: manager,
|
||||
deny_operation: false,
|
||||
denied_table: None,
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: label_metric_names.into_iter().map(String::from).collect(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn query_metric_name_values(
|
||||
handler: Arc<TestPrometheusHandler>,
|
||||
matches: Vec<&str>,
|
||||
) -> Vec<String> {
|
||||
let state: PrometheusHandlerRef = handler;
|
||||
let response = label_values_query(
|
||||
State(state),
|
||||
Path(METRIC_NAME_LABEL.to_string()),
|
||||
Extension(QueryContext::with(
|
||||
DEFAULT_CATALOG_NAME,
|
||||
DEFAULT_SCHEMA_NAME,
|
||||
)),
|
||||
Query(LabelValueQuery {
|
||||
matches: Matches(matches.into_iter().map(String::from).collect()),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
response.status_code.is_none(),
|
||||
"status={:?}, error={:?}",
|
||||
response.status_code,
|
||||
response.error
|
||||
);
|
||||
match response.data {
|
||||
PrometheusResponse::LabelValues(values) => values,
|
||||
other => panic!("expected label values, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn label_matchers_resolve_metric_names_from_data() {
|
||||
let handler = label_values_handler(vec!["cpu_user"]);
|
||||
let values = query_metric_name_values(handler.clone(), vec![r#"{pod="abc"}"#]).await;
|
||||
|
||||
// The metadata path would have enumerated both metrics.
|
||||
assert_eq!(vec!["cpu_user".to_string()], values);
|
||||
|
||||
let lookups = handler.label_lookups.lock().unwrap();
|
||||
assert_eq!(1, lookups.len());
|
||||
assert_eq!(
|
||||
vec!["pod".to_string()],
|
||||
lookups[0]
|
||||
.iter()
|
||||
.map(|matcher| matcher.name.clone())
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn metric_name_matchers_narrow_data_resolved_names() {
|
||||
let handler = label_values_handler(vec!["cpu_user", "cpu_system"]);
|
||||
let values =
|
||||
query_metric_name_values(handler, vec![r#"{__name__=~"cpu_u.*", pod="abc"}"#]).await;
|
||||
|
||||
assert_eq!(vec!["cpu_user".to_string()], values);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn special_matchers_are_stripped_before_the_data_lookup() {
|
||||
let handler = label_values_handler(vec!["cpu_user"]);
|
||||
let values = query_metric_name_values(
|
||||
handler.clone(),
|
||||
vec![r#"{pod="abc", __field__="value", __database__="public"}"#],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(vec!["cpu_user".to_string()], values);
|
||||
let lookups = handler.label_lookups.lock().unwrap();
|
||||
assert_eq!(
|
||||
vec!["pod".to_string()],
|
||||
lookups[0]
|
||||
.iter()
|
||||
.map(|matcher| matcher.name.clone())
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn database_and_field_matchers_stay_on_the_metadata_path() {
|
||||
let handler = label_values_handler(vec!["never_returned"]);
|
||||
let values = query_metric_name_values(
|
||||
handler.clone(),
|
||||
vec![r#"{__name__=~"cpu_.*", __field__="value", __database__="other"}"#],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
vec!["cpu_system".to_string(), "cpu_user".to_string()],
|
||||
values
|
||||
);
|
||||
assert!(handler.label_lookups.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_match_still_enumerates_every_metric() {
|
||||
let handler = label_values_handler(Vec::new());
|
||||
let values = query_metric_name_values(handler, Vec::new()).await;
|
||||
|
||||
assert_eq!(
|
||||
vec!["cpu_system".to_string(), "cpu_user".to_string()],
|
||||
values
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_series_query_expands_metric_name_regex() {
|
||||
let cpu_user = test_table_info(
|
||||
@@ -2728,6 +3116,8 @@ mod tests {
|
||||
deny_operation: false,
|
||||
denied_table: None,
|
||||
metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()],
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
});
|
||||
@@ -2889,6 +3279,134 @@ mod tests {
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
// Arithmetic operators always drop the metric name.
|
||||
TestCase {
|
||||
name: "arithmetic between metrics",
|
||||
promql: "a + b",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "arithmetic with scalar",
|
||||
promql: "a * 2",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
// Comparisons keep the metric name from the side whose labels the
|
||||
// result carries, i.e. the "many" side of the match: normally the
|
||||
// left operand, but `group_right` takes it from the right operand and
|
||||
// a scalar left operand leaves the vector on the right. The name is
|
||||
// dropped when the comparison returns a bool or is a one-to-one
|
||||
// `on(...)` match that does not keep `__name__` or an `ignoring(...)`
|
||||
// match that excludes it. A `group_left`/`group_right` modifier does
|
||||
// not reduce the labels, so the name is kept.
|
||||
TestCase {
|
||||
name: "bool comparison between metrics",
|
||||
promql: "a > bool b",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "bool comparison with scalar",
|
||||
promql: "a > bool 0.5",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with scalar",
|
||||
promql: "a > 0.5",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with scalar on the left",
|
||||
promql: "0.5 < a",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "equality comparison with scalar",
|
||||
promql: "up == 0",
|
||||
expected_metric: Some("up"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison between metrics",
|
||||
promql: "a > b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with ignoring",
|
||||
promql: "a > ignoring(x) b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with ignoring __name__",
|
||||
promql: "a > ignoring(__name__) b",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with ignoring __name__ and group_left",
|
||||
promql: "a > ignoring(__name__) group_left b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with ignoring __name__ and group_right",
|
||||
promql: "a > ignoring(__name__) group_right b",
|
||||
expected_metric: Some("b"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "unless with ignoring __name__",
|
||||
promql: "a unless ignoring(__name__) b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with on",
|
||||
promql: "a > on(x) b",
|
||||
expected_metric: None,
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with on keeping __name__",
|
||||
promql: "a > on(__name__) b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with on and group_left",
|
||||
promql: "a > on(x) group_left b",
|
||||
expected_metric: Some("a"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
TestCase {
|
||||
name: "comparison with on and group_right",
|
||||
promql: "a > on(x) group_right b",
|
||||
expected_metric: Some("b"),
|
||||
expected_type: ValueType::Vector,
|
||||
should_error: false,
|
||||
},
|
||||
// Subquery cases
|
||||
TestCase {
|
||||
name: "basic subquery",
|
||||
@@ -3635,6 +4153,8 @@ mod tests {
|
||||
deny_operation: true,
|
||||
denied_table: None,
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
})),
|
||||
@@ -3649,6 +4169,8 @@ mod tests {
|
||||
deny_operation: false,
|
||||
denied_table: Some("denied"),
|
||||
metric_names: Vec::new(),
|
||||
label_metric_names: Vec::new(),
|
||||
label_lookups: Mutex::new(Vec::new()),
|
||||
queries: Mutex::new(Vec::new()),
|
||||
ordered_outputs: Mutex::new(Vec::new()),
|
||||
});
|
||||
|
||||
@@ -30,6 +30,9 @@ pub enum TraceCoerceError {
|
||||
// - String to Int64
|
||||
// - String to Float64
|
||||
// - String to Boolean
|
||||
// The empty string and "-" coerce to null instead of failing: instrumentations
|
||||
// fill numeric attributes from absent headers this way (PSR-7 `getHeaderLine`
|
||||
// returns ""), and rejecting the value would reject the whole span.
|
||||
//
|
||||
// Lossless signed-to-unsigned integer casts. These let the built-in data
|
||||
// models move from unsigned to signed integers while existing unsigned tables
|
||||
@@ -86,6 +89,13 @@ pub fn coerce_value_data(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if let ValueData::StringValue(s) = v
|
||||
&& matches!(s.as_str(), "" | "-")
|
||||
&& is_supported_trace_coercion(request_type, target)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(value) = coerce_non_null_value(target, request_type, v) else {
|
||||
return Err(TraceCoerceError::Unsupported);
|
||||
};
|
||||
@@ -273,6 +283,24 @@ mod tests {
|
||||
assert_eq!(result, Err(TraceCoerceError::Unsupported));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_absent_value_strings_to_null() {
|
||||
for s in ["", "-"] {
|
||||
for target in [
|
||||
ColumnDataType::Int64,
|
||||
ColumnDataType::Float64,
|
||||
ColumnDataType::Boolean,
|
||||
] {
|
||||
let result = coerce_value_data(
|
||||
&Some(ValueData::StringValue(s.to_string())),
|
||||
target,
|
||||
ColumnDataType::String,
|
||||
);
|
||||
assert_eq!(result, Ok(None), "{s:?} to {target:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_float64_to_int64_not_supported() {
|
||||
let result = coerce_value_data(
|
||||
|
||||
+151
-23
@@ -362,22 +362,37 @@ pub fn recordbatches_to_timeseries(
|
||||
value_column_name: &str,
|
||||
recordbatches: RecordBatches,
|
||||
) -> Result<Vec<TimeSeries>> {
|
||||
Ok(recordbatches
|
||||
.take()
|
||||
.into_iter()
|
||||
.map(|x| recordbatch_to_timeseries(table_name, timestamp_column_name, value_column_name, x))
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect())
|
||||
let mut timeseries: Vec<TimeSeries> = Vec::new();
|
||||
let mut timeseries_by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
|
||||
|
||||
for recordbatch in recordbatches.take() {
|
||||
append_recordbatch_to_timeseries(
|
||||
table_name,
|
||||
timestamp_column_name,
|
||||
value_column_name,
|
||||
recordbatch,
|
||||
&mut timeseries,
|
||||
&mut timeseries_by_hash,
|
||||
)?;
|
||||
}
|
||||
|
||||
for ts in &mut timeseries {
|
||||
ts.samples.sort_unstable_by_key(|s| s.timestamp);
|
||||
}
|
||||
|
||||
timeseries
|
||||
.sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels));
|
||||
Ok(timeseries)
|
||||
}
|
||||
|
||||
fn recordbatch_to_timeseries(
|
||||
fn append_recordbatch_to_timeseries(
|
||||
table: &str,
|
||||
timestamp_column_name: &str,
|
||||
value_column_name: &str,
|
||||
recordbatch: RecordBatch,
|
||||
) -> Result<Vec<TimeSeries>> {
|
||||
timeseries: &mut Vec<TimeSeries>,
|
||||
timeseries_by_hash: &mut HashMap<u64, Vec<usize>>,
|
||||
) -> Result<()> {
|
||||
let ts_column = recordbatch
|
||||
.column_by_name(timestamp_column_name)
|
||||
.with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
|
||||
@@ -409,8 +424,6 @@ fn recordbatch_to_timeseries(
|
||||
})?;
|
||||
|
||||
let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?;
|
||||
let mut timeseries: Vec<TimeSeries> = Vec::new();
|
||||
let mut timeseries_by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
|
||||
let mut previous_timeseries: Option<usize> = None;
|
||||
|
||||
for row in 0..recordbatch.num_rows() {
|
||||
@@ -447,9 +460,7 @@ fn recordbatch_to_timeseries(
|
||||
timeseries[timeseries_index].samples.push(sample);
|
||||
}
|
||||
|
||||
timeseries
|
||||
.sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels));
|
||||
Ok(timeseries)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> {
|
||||
@@ -1122,6 +1133,121 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recordbatches_to_timeseries_merges_across_batches() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
greptime_timestamp(),
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
true,
|
||||
),
|
||||
ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
|
||||
ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
|
||||
]));
|
||||
|
||||
let recordbatches = RecordBatches::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
RecordBatch::new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondVector::from_vec(vec![3000, 1500])) as _,
|
||||
Arc::new(Float64Vector::from_vec(vec![30.0, 15.0])) as _,
|
||||
Arc::new(StringVector::from(vec!["host1", "host2"])) as _,
|
||||
],
|
||||
)
|
||||
.unwrap(),
|
||||
RecordBatch::new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _,
|
||||
Arc::new(Float64Vector::from_vec(vec![10.0])) as _,
|
||||
Arc::new(StringVector::from(vec!["host1"])) as _,
|
||||
],
|
||||
)
|
||||
.unwrap(),
|
||||
RecordBatch::new(
|
||||
schema,
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondVector::from_vec(vec![2000, 2500])) as _,
|
||||
Arc::new(Float64Vector::from_vec(vec![20.0, 25.0])) as _,
|
||||
Arc::new(StringVector::from(vec!["host1", "host2"])) as _,
|
||||
],
|
||||
)
|
||||
.unwrap(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let timeseries = recordbatches_to_timeseries(
|
||||
"cpu_usage",
|
||||
greptime_timestamp(),
|
||||
greptime_value(),
|
||||
recordbatches,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(2, timeseries.len());
|
||||
|
||||
assert_eq!(
|
||||
vec![
|
||||
Label {
|
||||
name: METRIC_NAME_LABEL.to_string(),
|
||||
value: "cpu_usage".to_string(),
|
||||
},
|
||||
Label {
|
||||
name: "instance".to_string(),
|
||||
value: "host1".to_string(),
|
||||
},
|
||||
],
|
||||
timeseries[0].labels
|
||||
);
|
||||
assert_eq!(
|
||||
vec![
|
||||
Sample {
|
||||
value: 10.0,
|
||||
timestamp: 1000,
|
||||
},
|
||||
Sample {
|
||||
value: 20.0,
|
||||
timestamp: 2000,
|
||||
},
|
||||
Sample {
|
||||
value: 30.0,
|
||||
timestamp: 3000,
|
||||
},
|
||||
],
|
||||
timeseries[0].samples
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
vec![
|
||||
Label {
|
||||
name: METRIC_NAME_LABEL.to_string(),
|
||||
value: "cpu_usage".to_string(),
|
||||
},
|
||||
Label {
|
||||
name: "instance".to_string(),
|
||||
value: "host2".to_string(),
|
||||
},
|
||||
],
|
||||
timeseries[1].labels
|
||||
);
|
||||
assert_eq!(
|
||||
vec![
|
||||
Sample {
|
||||
value: 15.0,
|
||||
timestamp: 1500,
|
||||
},
|
||||
Sample {
|
||||
value: 25.0,
|
||||
timestamp: 2500,
|
||||
},
|
||||
],
|
||||
timeseries[1].samples
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recordbatches_to_timeseries_borrows_and_groups_dictionary_labels() {
|
||||
let arrow_schema = Arc::new(ArrowSchema::new(vec![
|
||||
@@ -1212,7 +1338,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recordbatch_to_timeseries_groups_non_contiguous_series() {
|
||||
fn test_recordbatches_to_timeseries_groups_non_contiguous_series() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
greptime_timestamp(),
|
||||
@@ -1223,7 +1349,7 @@ mod tests {
|
||||
ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
|
||||
]));
|
||||
let recordbatch = RecordBatch::new(
|
||||
schema,
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _,
|
||||
Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _,
|
||||
@@ -1232,11 +1358,12 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let timeseries = recordbatch_to_timeseries(
|
||||
let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap();
|
||||
let timeseries = recordbatches_to_timeseries(
|
||||
"metric1",
|
||||
greptime_timestamp(),
|
||||
greptime_value(),
|
||||
recordbatch,
|
||||
recordbatches,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
@@ -1259,7 +1386,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recordbatch_to_timeseries_arrow_label_types_and_nulls() {
|
||||
fn test_recordbatches_to_timeseries_arrow_label_types_and_nulls() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
greptime_timestamp(),
|
||||
@@ -1272,7 +1399,7 @@ mod tests {
|
||||
ColumnSchema::new("shard", ConcreteDataType::int32_datatype(), true),
|
||||
]));
|
||||
let recordbatch = RecordBatch::new(
|
||||
schema,
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _,
|
||||
Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _,
|
||||
@@ -1289,11 +1416,12 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let timeseries = recordbatch_to_timeseries(
|
||||
let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap();
|
||||
let timeseries = recordbatches_to_timeseries(
|
||||
"metric1",
|
||||
greptime_timestamp(),
|
||||
greptime_value(),
|
||||
recordbatch,
|
||||
recordbatches,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -184,6 +184,21 @@ pub trait PrometheusHandler {
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<String>>;
|
||||
|
||||
/// Query metric table names that carry data matching `matchers` in the time
|
||||
/// range. `matchers` must hold only ordinary label matchers: `__name__`
|
||||
/// names a table and the database and field matchers name no column, so the
|
||||
/// caller resolves all three before calling.
|
||||
///
|
||||
/// Only metric engine tables are covered.
|
||||
async fn query_metric_names_by_labels(
|
||||
&self,
|
||||
matchers: Vec<Matcher>,
|
||||
schema: &str,
|
||||
start: SystemTime,
|
||||
end: SystemTime,
|
||||
ctx: &QueryContextRef,
|
||||
) -> Result<Vec<String>>;
|
||||
|
||||
async fn query_label_values(
|
||||
&self,
|
||||
metric: String,
|
||||
|
||||
@@ -704,6 +704,125 @@ LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' AND k < '8'
|
||||
}
|
||||
}
|
||||
|
||||
/// Regression test for `MatchesConstantTermOptimizer` dropping `FilterExec`'s `fetch`.
|
||||
///
|
||||
/// On a single-partition plan the `LIMIT` is enforced by the `fetch` that the physical
|
||||
/// `LimitPushdown` rule leaves on the `FilterExec`, so a rule rebuilding that `FilterExec`
|
||||
/// has to carry the `fetch` over. While it does not, the limit is silently lifted for
|
||||
/// `matches_term` filters and
|
||||
/// `SELECT count(*) FROM (SELECT ts, msg FROM t WHERE matches_term(msg, 'foo') LIMIT 1) s`
|
||||
/// returns the number of matching rows (2) instead of 1.
|
||||
///
|
||||
/// `query_parallelism=4` is expected to stay correct, because the limit is then materialized
|
||||
/// above the scan instead of being delegated to the `FilterExec`; the test pins down that
|
||||
/// contrast. The metric-engine logical table covers the same filter on a shared physical
|
||||
/// region.
|
||||
///
|
||||
/// Both matching rows share `ts` and `msg`, so the limited output does not depend on the
|
||||
/// order in which the scan returns them.
|
||||
#[apply(both_instances_cases)]
|
||||
async fn test_matches_term_subquery_limit_count(instance: Arc<dyn MockInstance>) {
|
||||
let frontend = instance.frontend();
|
||||
|
||||
execute_sql(
|
||||
&frontend,
|
||||
r#"
|
||||
CREATE TABLE f07_mito_msgs (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
host STRING PRIMARY KEY,
|
||||
msg STRING,
|
||||
val DOUBLE
|
||||
)
|
||||
ENGINE = mito
|
||||
"#,
|
||||
)
|
||||
.await;
|
||||
execute_sql(
|
||||
&frontend,
|
||||
r#"INSERT INTO f07_mito_msgs (ts, host, msg, val) VALUES
|
||||
(0, 'host1', 'foo bar', 1.0),
|
||||
(0, 'host2', 'foo bar', 2.0),
|
||||
(0, 'host3', 'other baz', 3.0)"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
execute_sql(
|
||||
&frontend,
|
||||
r#"
|
||||
CREATE TABLE f07_metric_phy (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
val DOUBLE,
|
||||
host STRING,
|
||||
msg STRING,
|
||||
PRIMARY KEY (host, msg)
|
||||
)
|
||||
ENGINE = metric
|
||||
WITH ('physical_metric_table' = '')
|
||||
"#,
|
||||
)
|
||||
.await;
|
||||
execute_sql(
|
||||
&frontend,
|
||||
r#"
|
||||
CREATE TABLE f07_metric_msgs (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
val DOUBLE,
|
||||
host STRING,
|
||||
msg STRING,
|
||||
PRIMARY KEY (host, msg)
|
||||
)
|
||||
ENGINE = metric
|
||||
WITH ('on_physical_table' = 'f07_metric_phy')
|
||||
"#,
|
||||
)
|
||||
.await;
|
||||
execute_sql(
|
||||
&frontend,
|
||||
r#"INSERT INTO f07_metric_msgs (ts, host, msg, val) VALUES
|
||||
(0, 'host1', 'foo bar', 1.0),
|
||||
(0, 'host2', 'foo bar', 2.0),
|
||||
(0, 'host3', 'other baz', 3.0)"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
for table in ["f07_mito_msgs", "f07_metric_msgs"] {
|
||||
let limited = format!("SELECT ts, msg FROM {table} WHERE matches_term(msg, 'foo') LIMIT 1");
|
||||
let count = format!("SELECT count(*) FROM ({limited}) s");
|
||||
|
||||
let limit_rows = execute_sql_with_query_parallelism(&frontend, &limited, 1)
|
||||
.await
|
||||
.data
|
||||
.pretty_print()
|
||||
.await;
|
||||
assert_eq!(
|
||||
limit_rows,
|
||||
r#"+---------------------+---------+
|
||||
| ts | msg |
|
||||
+---------------------+---------+
|
||||
| 1970-01-01T00:00:00 | foo bar |
|
||||
+---------------------+---------+"#,
|
||||
"query_parallelism=1: `{limited}` must return exactly the single limited row"
|
||||
);
|
||||
|
||||
for parallelism in [1, 4] {
|
||||
let count_rows = execute_sql_with_query_parallelism(&frontend, &count, parallelism)
|
||||
.await
|
||||
.data
|
||||
.pretty_print()
|
||||
.await;
|
||||
assert_eq!(
|
||||
count_rows,
|
||||
r#"+----------+
|
||||
| count(*) |
|
||||
+----------+
|
||||
| 1 |
|
||||
+----------+"#,
|
||||
"query_parallelism={parallelism}: the LIMIT 1 of the subquery must cap `{count}` to one row"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[apply(both_instances_cases)]
|
||||
async fn test_partitioned_merge_scan_semi_join_repartitions(instance: Arc<dyn MockInstance>) {
|
||||
let frontend = instance.frontend();
|
||||
|
||||
@@ -1516,6 +1516,117 @@ pub async fn test_prom_http_api(store_type: StorageType) {
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
// query `__name__` by a matcher on an ordinary label: the metric engine
|
||||
// physical tables are scanned, so only metrics carrying the label value are
|
||||
// returned. `demo_metrics` shares `phy` with `demo` but has no `host` value.
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=0&end=600")
|
||||
.send()
|
||||
.await;
|
||||
let status = res.status();
|
||||
let text = res.text().await;
|
||||
assert_eq!(status, StatusCode::OK, "{text}");
|
||||
let prom_resp = serde_json::from_str::<PrometheusJsonResponse>(&text).unwrap();
|
||||
assert_eq!(prom_resp.status, "success");
|
||||
assert!(prom_resp.error.is_none());
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!(["demo", "multi_labels"])).unwrap()
|
||||
);
|
||||
|
||||
// `__name__` matchers narrow the names the data resolved: `multi_labels`
|
||||
// also carries `idc="idc1"` but its name does not match.
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\"demo.*\", idc=\"idc1\"}&start=0&end=600")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(prom_resp.status, "success");
|
||||
assert!(prom_resp.error.is_none());
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!([
|
||||
"demo_metrics",
|
||||
"demo_metrics_with_nanos",
|
||||
]))
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
// The time range selects the series: `demo` carries `host="host2"` only at
|
||||
// t=600, so narrowing the range drops it while `multi_labels` at t=0 stays.
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=600")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!(["demo", "multi_labels"])).unwrap()
|
||||
);
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host2\"}&start=0&end=300")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!(["multi_labels"])).unwrap()
|
||||
);
|
||||
|
||||
// Logical metrics sharing a physical table have NULL in the label columns
|
||||
// they don't use. Prometheus reads a label a series doesn't carry as the
|
||||
// empty string, so `demo_metrics` and `demo_metrics_with_nanos` — neither of
|
||||
// which has a `host` label — match both of these.
|
||||
//
|
||||
// `.%2B` is `.+`; a bare `+` decodes to a space in a query string. Grafana
|
||||
// encodes it the same way.
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host=\"\"}&start=0&end=600")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!([
|
||||
"demo_metrics",
|
||||
"demo_metrics_with_nanos",
|
||||
]))
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={__name__=~\".%2B\", host!=\"host1\"}&start=0&end=600")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!([
|
||||
"demo",
|
||||
"demo_metrics",
|
||||
"demo_metrics_with_nanos",
|
||||
"multi_labels",
|
||||
]))
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
// A pre-epoch RFC3339 bound is a valid range, not a panic.
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/label/__name__/values?match[]={host=\"host1\"}&start=1969-12-31T23:59:59Z&end=600")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
let prom_resp = res.json::<PrometheusJsonResponse>().await;
|
||||
assert_eq!(
|
||||
prom_resp.data,
|
||||
serde_json::from_value::<PrometheusResponse>(json!(["demo", "multi_labels"])).unwrap()
|
||||
);
|
||||
|
||||
// buildinfo
|
||||
let res = client
|
||||
.get("/v1/prometheus/api/v1/status/buildinfo")
|
||||
|
||||
@@ -144,3 +144,94 @@ DROP TABLE doubles;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
-- On partitioned tables the aggregate is split into partial state and merge
|
||||
CREATE TABLE array_agg_partitioned (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
k INT,
|
||||
lat DOUBLE,
|
||||
PRIMARY KEY(k)
|
||||
)
|
||||
PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO array_agg_partitioned VALUES
|
||||
(1000, 1, 1),
|
||||
(2000, 11, 2),
|
||||
(3000, 21, 3),
|
||||
(4000, 2, 4),
|
||||
(5000, 12, 5),
|
||||
(6000, 22, 6),
|
||||
(7000, 3, 7);
|
||||
|
||||
Affected Rows: 7
|
||||
|
||||
SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned;
|
||||
|
||||
+-----------------------------------------------------------------------------------------+
|
||||
| array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] |
|
||||
+-----------------------------------------------------------------------------------------+
|
||||
| [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] |
|
||||
+-----------------------------------------------------------------------------------------+
|
||||
|
||||
SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
+-------+-------------------------------------------------------------------------------------------+
|
||||
| g | array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] |
|
||||
+-------+-------------------------------------------------------------------------------------------+
|
||||
| false | [3.0, 2.0, 1.0] |
|
||||
| true | [7.0, 6.0, 5.0, 4.0] |
|
||||
+-------+-------------------------------------------------------------------------------------------+
|
||||
|
||||
SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned;
|
||||
|
||||
+---------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| array_agg(array_agg_partitioned.k) ORDER BY [array_agg_partitioned.k % Int64(10) ASC NULLS LAST, array_agg_partitioned.ts DESC NULLS FIRST] |
|
||||
+---------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| [21, 11, 1, 22, 12, 2, 3] |
|
||||
+---------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
|
||||
-- nth_value needs sorted input, so it isn't split into partial state and merge
|
||||
SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned;
|
||||
|
||||
+--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+
|
||||
| nth_value(array_agg_partitioned.lat,Int64(2)) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | nth_value(array_agg_partitioned.lat,Int64(3)) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] |
|
||||
+--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+
|
||||
| 2.0 | 5.0 |
|
||||
+--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+
|
||||
|
||||
-- WITHIN GROUP aggregates don't need sorted input and are still split
|
||||
-- SQLNESS REPLACE (peers.*) REDACTED
|
||||
-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
|
||||
-- SQLNESS REPLACE (-+) -
|
||||
-- SQLNESS REPLACE (\s\s+) _
|
||||
EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned;
|
||||
|
||||
+-+-+
|
||||
| plan_type_| plan_|
|
||||
+-+-+
|
||||
| logical_plan_| Aggregate: groupBy=[[]], aggr=[[__sum_merge(__sum_state(array_agg_partitioned.lat)) AS sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) AS approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]] |
|
||||
|_|_MergeScan [is_placeholder=false, remote_input=[_|
|
||||
|_| Aggregate: groupBy=[[]], aggr=[[__sum_state(array_agg_partitioned.lat), __approx_percentile_cont_state(array_agg_partitioned.lat, Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]]]_|
|
||||
|_|_TableScan: array_agg_partitioned_|
|
||||
|_| ]]_|
|
||||
| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[sum(array_agg_partitioned.lat), approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_|
|
||||
|_|_CoalescePartitionsExec_|
|
||||
|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(array_agg_partitioned.lat), approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_|
|
||||
|_|_RepartitionExec: partitioning=REDACTED
|
||||
|_|_MergeScanExec: REDACTED
|
||||
|_|_|
|
||||
+-+-+
|
||||
|
||||
SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned;
|
||||
|
||||
+--------------------------------+-------------------------------------------------------------------------------------------------+
|
||||
| sum(array_agg_partitioned.lat) | approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST] |
|
||||
+--------------------------------+-------------------------------------------------------------------------------------------------+
|
||||
| 28.0 | 2.25 |
|
||||
+--------------------------------+-------------------------------------------------------------------------------------------------+
|
||||
|
||||
DROP TABLE array_agg_partitioned;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
|
||||
@@ -54,3 +54,41 @@ DROP TABLE integers;
|
||||
DROP TABLE strings;
|
||||
|
||||
DROP TABLE doubles;
|
||||
|
||||
-- On partitioned tables the aggregate is split into partial state and merge
|
||||
CREATE TABLE array_agg_partitioned (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
k INT,
|
||||
lat DOUBLE,
|
||||
PRIMARY KEY(k)
|
||||
)
|
||||
PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20);
|
||||
|
||||
INSERT INTO array_agg_partitioned VALUES
|
||||
(1000, 1, 1),
|
||||
(2000, 11, 2),
|
||||
(3000, 21, 3),
|
||||
(4000, 2, 4),
|
||||
(5000, 12, 5),
|
||||
(6000, 22, 6),
|
||||
(7000, 3, 7);
|
||||
|
||||
SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned;
|
||||
|
||||
SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned;
|
||||
|
||||
-- nth_value needs sorted input, so it isn't split into partial state and merge
|
||||
SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned;
|
||||
|
||||
-- WITHIN GROUP aggregates don't need sorted input and are still split
|
||||
-- SQLNESS REPLACE (peers.*) REDACTED
|
||||
-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
|
||||
-- SQLNESS REPLACE (-+) -
|
||||
-- SQLNESS REPLACE (\s\s+) _
|
||||
EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned;
|
||||
|
||||
SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned;
|
||||
|
||||
DROP TABLE array_agg_partitioned;
|
||||
|
||||
@@ -437,3 +437,47 @@ FROM
|
||||
| true | false | true | false | false | true |
|
||||
+--------------------------+--------------------------+------------------------+------------------------+----------------------------------+----------------------------------+
|
||||
|
||||
CREATE TABLE geo_path_partitioned (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
k INT,
|
||||
lat DOUBLE,
|
||||
lon DOUBLE,
|
||||
PRIMARY KEY(k)
|
||||
)
|
||||
PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO geo_path_partitioned VALUES
|
||||
(1000, 1, 1, 11),
|
||||
(2000, 11, 2, 12),
|
||||
(3000, 21, 3, 13),
|
||||
(4000, 2, 4, 14),
|
||||
(5000, 12, 5, 15),
|
||||
(6000, 22, 6, 16),
|
||||
(7000, 3, 7, 17);
|
||||
|
||||
Affected Rows: 7
|
||||
|
||||
SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
+-------+-------------------------------------------------------------------------------------+
|
||||
| g | geo_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) |
|
||||
+-------+-------------------------------------------------------------------------------------+
|
||||
| false | {lat: [1.0, 2.0, 3.0], lng: [11.0, 12.0, 13.0]} |
|
||||
| true | {lat: [4.0, 5.0, 6.0, 7.0], lng: [14.0, 15.0, 16.0, 17.0]} |
|
||||
+-------+-------------------------------------------------------------------------------------+
|
||||
|
||||
SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
+-------+---------------------------------------------------------------------------------------------+
|
||||
| g | json_encode_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) |
|
||||
+-------+---------------------------------------------------------------------------------------------+
|
||||
| false | [[11.0,1.0],[12.0,2.0],[13.0,3.0]] |
|
||||
| true | [[14.0,4.0],[15.0,5.0],[16.0,6.0],[17.0,7.0]] |
|
||||
+-------+---------------------------------------------------------------------------------------------+
|
||||
|
||||
DROP TABLE geo_path_partitioned;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
|
||||
@@ -180,3 +180,27 @@ FROM
|
||||
'POLYGON ((-121.491698 38.653343, -121.582353 38.556757, -121.469721 38.449287, -121.315883 38.541721, -121.491698 38.653343))' AS polygon2,
|
||||
'POLYGON ((-122.089628 37.450332, -122.20535 37.378342, -122.093062 37.36088, -122.044301 37.372886, -122.089628 37.450332))' AS polygon3,
|
||||
);
|
||||
|
||||
CREATE TABLE geo_path_partitioned (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
k INT,
|
||||
lat DOUBLE,
|
||||
lon DOUBLE,
|
||||
PRIMARY KEY(k)
|
||||
)
|
||||
PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20);
|
||||
|
||||
INSERT INTO geo_path_partitioned VALUES
|
||||
(1000, 1, 1, 11),
|
||||
(2000, 11, 2, 12),
|
||||
(3000, 21, 3, 13),
|
||||
(4000, 2, 4, 14),
|
||||
(5000, 12, 5, 15),
|
||||
(6000, 22, 6, 16),
|
||||
(7000, 3, 7, 17);
|
||||
|
||||
SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g;
|
||||
|
||||
DROP TABLE geo_path_partitioned;
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
-- A vector combined with a computed scalar on the right keeps the vector's
|
||||
-- labels, like a scalar on the left does. Expected results are from Prometheus 3.
|
||||
create table computed_scalar_lhs (
|
||||
ts timestamp time index,
|
||||
val double,
|
||||
host string primary key
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
create table computed_scalar_rhs (
|
||||
ts timestamp time index,
|
||||
val double,
|
||||
host string primary key
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
insert into computed_scalar_lhs values
|
||||
(0, 1.0, 'a'),
|
||||
(0, 2.0, 'b');
|
||||
|
||||
Affected Rows: 2
|
||||
|
||||
insert into computed_scalar_rhs values
|
||||
(0, 10.0, 'x');
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs + time();
|
||||
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
| host | ts | computed_scalar_lhs.val + .time / Float64(1000) |
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 61.0 |
|
||||
| b | 1970-01-01T00:01:00 | 62.0 |
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') time() + computed_scalar_lhs;
|
||||
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
| host | ts | .time / Float64(1000) + computed_scalar_lhs.val |
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 61.0 |
|
||||
| b | 1970-01-01T00:01:00 | 62.0 |
|
||||
+------+---------------------+-------------------------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs);
|
||||
|
||||
+------+---------------------+-----------------------------------------------------------+
|
||||
| host | ts | computed_scalar_lhs.val * computed_scalar_rhs.scalar(val) |
|
||||
+------+---------------------+-----------------------------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 10.0 |
|
||||
| b | 1970-01-01T00:01:00 | 20.0 |
|
||||
+------+---------------------+-----------------------------------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2);
|
||||
|
||||
+------+---------------------+--------------------------------------------------------------+
|
||||
| host | ts | computed_scalar_lhs.val < .time / Float64(1000) * Float64(2) |
|
||||
+------+---------------------+--------------------------------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 1.0 |
|
||||
| b | 1970-01-01T00:01:00 | 1.0 |
|
||||
+------+---------------------+--------------------------------------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs);
|
||||
|
||||
+------+---------------------+---------------------------------------------+
|
||||
| host | ts | power(.val,computed_scalar_rhs.scalar(val)) |
|
||||
+------+---------------------+---------------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 1.0 |
|
||||
| b | 1970-01-01T00:01:00 | 1024.0 |
|
||||
+------+---------------------+---------------------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs);
|
||||
|
||||
+------+---------------------+----------------------------------------+
|
||||
| host | ts | computed_scalar_rhs.scalar(val) * .val |
|
||||
+------+---------------------+----------------------------------------+
|
||||
| a | 1970-01-01T00:01:00 | 10.0 |
|
||||
| b | 1970-01-01T00:01:00 | 20.0 |
|
||||
+------+---------------------+----------------------------------------+
|
||||
|
||||
drop table computed_scalar_lhs;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
drop table computed_scalar_rhs;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
-- A vector combined with a computed scalar on the right keeps the vector's
|
||||
-- labels, like a scalar on the left does. Expected results are from Prometheus 3.
|
||||
|
||||
create table computed_scalar_lhs (
|
||||
ts timestamp time index,
|
||||
val double,
|
||||
host string primary key
|
||||
);
|
||||
|
||||
create table computed_scalar_rhs (
|
||||
ts timestamp time index,
|
||||
val double,
|
||||
host string primary key
|
||||
);
|
||||
|
||||
insert into computed_scalar_lhs values
|
||||
(0, 1.0, 'a'),
|
||||
(0, 2.0, 'b');
|
||||
|
||||
insert into computed_scalar_rhs values
|
||||
(0, 10.0, 'x');
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs + time();
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') time() + computed_scalar_lhs;
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs);
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2);
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs);
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs);
|
||||
|
||||
drop table computed_scalar_lhs;
|
||||
|
||||
drop table computed_scalar_rhs;
|
||||
@@ -55,15 +55,15 @@ TQL EVAL (0, 20, '5s') pi() + 1;
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
TQL EVAL (0, 20, '5s') angles{unit="radians"} * pi();
|
||||
|
||||
+---------------------+--------------------+
|
||||
| time | angles.val * .pi() |
|
||||
+---------------------+--------------------+
|
||||
| 1970-01-01T00:00:00 | 0.0 |
|
||||
| 1970-01-01T00:00:05 | 4.934813740258847 |
|
||||
| 1970-01-01T00:00:10 | 9.869627480517694 |
|
||||
| 1970-01-01T00:00:15 | 14.80444122077654 |
|
||||
| 1970-01-01T00:00:20 | 19.73925496103539 |
|
||||
+---------------------+--------------------+
|
||||
+---------+---------------------+--------------------+
|
||||
| unit | ts | angles.val * .pi() |
|
||||
+---------+---------------------+--------------------+
|
||||
| radians | 1970-01-01T00:00:00 | 0.0 |
|
||||
| radians | 1970-01-01T00:00:05 | 4.934813740258847 |
|
||||
| radians | 1970-01-01T00:00:10 | 9.869627480517694 |
|
||||
| radians | 1970-01-01T00:00:15 | 14.80444122077654 |
|
||||
| radians | 1970-01-01T00:00:20 | 19.73925496103539 |
|
||||
+---------+---------------------+--------------------+
|
||||
|
||||
-- Test rad() function (converts degrees to radians)
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
@@ -310,18 +310,18 @@ TQL EVAL (25, 35, '5s') sgn(angles - 42.5);
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
TQL EVAL (0, 35, '5s') sgn(angles{unit="radians"} - pi());
|
||||
|
||||
+---------------------+----------------------------+
|
||||
| time | signum(angles.val - .pi()) |
|
||||
+---------------------+----------------------------+
|
||||
| 1970-01-01T00:00:00 | -1.0 |
|
||||
| 1970-01-01T00:00:05 | -1.0 |
|
||||
| 1970-01-01T00:00:10 | 1.0 |
|
||||
| 1970-01-01T00:00:15 | 1.0 |
|
||||
| 1970-01-01T00:00:20 | 1.0 |
|
||||
| 1970-01-01T00:00:25 | 1.0 |
|
||||
| 1970-01-01T00:00:30 | 1.0 |
|
||||
| 1970-01-01T00:00:35 | 1.0 |
|
||||
+---------------------+----------------------------+
|
||||
+---------------------+----------------------------+---------+
|
||||
| ts | signum(angles.val - .pi()) | unit |
|
||||
+---------------------+----------------------------+---------+
|
||||
| 1970-01-01T00:00:00 | -1.0 | radians |
|
||||
| 1970-01-01T00:00:05 | -1.0 | radians |
|
||||
| 1970-01-01T00:00:10 | 1.0 | radians |
|
||||
| 1970-01-01T00:00:15 | 1.0 | radians |
|
||||
| 1970-01-01T00:00:20 | 1.0 | radians |
|
||||
| 1970-01-01T00:00:25 | 1.0 | radians |
|
||||
| 1970-01-01T00:00:30 | 1.0 | radians |
|
||||
| 1970-01-01T00:00:35 | 1.0 | radians |
|
||||
+---------------------+----------------------------+---------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
TQL EVAL (25, 35, '5s') sgn(angles) * angles;
|
||||
|
||||
@@ -144,8 +144,9 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"};
|
||||
|_|_|_|
|
||||
| 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED
|
||||
|_|_|_PromSeriesDivideExec: tags=["host"] REDACTED
|
||||
|_|_|_FilterExec: CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END = REDACTED
|
||||
|_|_|_CooperativeExec REDACTED
|
||||
|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED
|
||||
|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["CASE WHEN host IS NOT NULL THEN host ELSE Dictionary(UInt32, Utf8(\"\")) END = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED
|
||||
|_|_|_|
|
||||
|_|_| Total rows: 0_|
|
||||
+-+-+-+
|
||||
|
||||
@@ -428,3 +428,92 @@ drop table histogram5_bucket;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
-- Issue #9443: string labels stay distinct during aggregation, but equal
|
||||
-- numeric bucket boundaries must be coalesced before monotonicity repair.
|
||||
CREATE TABLE equivalent_bounds_bucket (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
instance STRING,
|
||||
le STRING,
|
||||
val DOUBLE,
|
||||
PRIMARY KEY (instance, le)
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO equivalent_bounds_bucket VALUES
|
||||
(3000000, 'a', '10', 40),
|
||||
(3000000, 'a', '100', 50),
|
||||
(3000000, 'a', '+Inf', 50),
|
||||
(3000000, 'b', '10.0', 40),
|
||||
(3000000, 'b', '100.0', 50),
|
||||
(3000000, 'b', '1000.0', 50),
|
||||
(3000000, 'b', '+Inf', 50);
|
||||
|
||||
Affected Rows: 7
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket);
|
||||
|
||||
+--------+---------------------+-----------------------------------+
|
||||
| le | ts | sum(equivalent_bounds_bucket.val) |
|
||||
+--------+---------------------+-----------------------------------+
|
||||
| +Inf | 1970-01-01T00:50:00 | 100.0 |
|
||||
| 10 | 1970-01-01T00:50:00 | 40.0 |
|
||||
| 10.0 | 1970-01-01T00:50:00 | 40.0 |
|
||||
| 100 | 1970-01-01T00:50:00 | 50.0 |
|
||||
| 100.0 | 1970-01-01T00:50:00 | 50.0 |
|
||||
| 1000.0 | 1970-01-01T00:50:00 | 50.0 |
|
||||
+--------+---------------------+-----------------------------------+
|
||||
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket));
|
||||
|
||||
+---------------------+-----------------------------------+
|
||||
| ts | sum(equivalent_bounds_bucket.val) |
|
||||
+---------------------+-----------------------------------+
|
||||
| 1970-01-01T00:50:00 | 77.5 |
|
||||
+---------------------+-----------------------------------+
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket));
|
||||
|
||||
+----------+---------------------+-----------------------------------+
|
||||
| instance | ts | sum(equivalent_bounds_bucket.val) |
|
||||
+----------+---------------------+-----------------------------------+
|
||||
| a | 1970-01-01T00:50:00 | 77.5 |
|
||||
| b | 1970-01-01T00:50:00 | 77.5 |
|
||||
+----------+---------------------+-----------------------------------+
|
||||
|
||||
DROP TABLE equivalent_bounds_bucket;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
-- Equivalent infinity labels belong to the same histogram until the group ends.
|
||||
CREATE TABLE infinity_bounds_bucket (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
le STRING,
|
||||
val DOUBLE,
|
||||
PRIMARY KEY (le)
|
||||
);
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
INSERT INTO infinity_bounds_bucket VALUES
|
||||
(3000000, '10', 20),
|
||||
(3000000, '10.0', 20),
|
||||
(3000000, '+Inf', 25),
|
||||
(3000000, '+inf', 25);
|
||||
|
||||
Affected Rows: 4
|
||||
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket);
|
||||
|
||||
+---------------------+------+
|
||||
| ts | val |
|
||||
+---------------------+------+
|
||||
| 1970-01-01T00:50:00 | 6.25 |
|
||||
+---------------------+------+
|
||||
|
||||
DROP TABLE infinity_bounds_bucket;
|
||||
|
||||
Affected Rows: 0
|
||||
|
||||
|
||||
@@ -253,3 +253,50 @@ insert into histogram5_bucket values
|
||||
tql eval (3000, 3015, '3s') histogram_quantile(0.5, histogram5_bucket);
|
||||
|
||||
drop table histogram5_bucket;
|
||||
|
||||
-- Issue #9443: string labels stay distinct during aggregation, but equal
|
||||
-- numeric bucket boundaries must be coalesced before monotonicity repair.
|
||||
CREATE TABLE equivalent_bounds_bucket (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
instance STRING,
|
||||
le STRING,
|
||||
val DOUBLE,
|
||||
PRIMARY KEY (instance, le)
|
||||
);
|
||||
|
||||
INSERT INTO equivalent_bounds_bucket VALUES
|
||||
(3000000, 'a', '10', 40),
|
||||
(3000000, 'a', '100', 50),
|
||||
(3000000, 'a', '+Inf', 50),
|
||||
(3000000, 'b', '10.0', 40),
|
||||
(3000000, 'b', '100.0', 50),
|
||||
(3000000, 'b', '1000.0', 50),
|
||||
(3000000, 'b', '+Inf', 50);
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket);
|
||||
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket));
|
||||
|
||||
-- SQLNESS SORT_RESULT 3 1
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket));
|
||||
|
||||
DROP TABLE equivalent_bounds_bucket;
|
||||
|
||||
-- Equivalent infinity labels belong to the same histogram until the group ends.
|
||||
CREATE TABLE infinity_bounds_bucket (
|
||||
ts TIMESTAMP TIME INDEX,
|
||||
le STRING,
|
||||
val DOUBLE,
|
||||
PRIMARY KEY (le)
|
||||
);
|
||||
|
||||
INSERT INTO infinity_bounds_bucket VALUES
|
||||
(3000000, '10', 20),
|
||||
(3000000, '10.0', 20),
|
||||
(3000000, '+Inf', 25),
|
||||
(3000000, '+inf', 25);
|
||||
|
||||
tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket);
|
||||
|
||||
DROP TABLE infinity_bounds_bucket;
|
||||
|
||||
@@ -150,7 +150,7 @@ tql eval (1, 2, '1s') time() == bool metrics;
|
||||
tql eval (1, 2, '1s') metrics + time();
|
||||
|
||||
+---------------------+-------------------------------------+
|
||||
| time | metrics.val + .time / Float64(1000) |
|
||||
| ts | metrics.val + .time / Float64(1000) |
|
||||
+---------------------+-------------------------------------+
|
||||
| 1970-01-01T00:00:01 | 2.0 |
|
||||
| 1970-01-01T00:00:02 | 4.0 |
|
||||
@@ -170,7 +170,7 @@ tql eval (1, 2, '1s') metrics == time();
|
||||
tql eval (1, 2, '1s') metrics == bool time();
|
||||
|
||||
+---------------------+-------------------------------------+
|
||||
| time | metrics.val = .time / Float64(1000) |
|
||||
| ts | metrics.val = .time / Float64(1000) |
|
||||
+---------------------+-------------------------------------+
|
||||
| 1970-01-01T00:00:01 | 1.0 |
|
||||
| 1970-01-01T00:00:02 | 1.0 |
|
||||
|
||||
@@ -95,7 +95,7 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) > bool 30;
|
||||
tql eval (0, 60, '30s') timestamp(timestamp_test) - time();
|
||||
|
||||
+---------------------+----------------------------------------------+
|
||||
| time | timestamp_test.value - .time / Float64(1000) |
|
||||
| ts | timestamp_test.value - .time / Float64(1000) |
|
||||
+---------------------+----------------------------------------------+
|
||||
| 1970-01-01T00:00:00 | 0.0 |
|
||||
| 1970-01-01T00:00:30 | -29.0 |
|
||||
|
||||
Reference in New Issue
Block a user