mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
perf(mito2): retain batch-backed labels for native samples (#10134)
Retain compacted metadata rows for new series candidates and materialize owned labels lazily. Preserve region-qualified identities, full physical-tag conflict checks and mixed-predicate fallback. Separate scan bindings from output labels and validate execution-time schemas. Validated cross-bucket retention and conflicts, filtered PromQL label binding, affected crate tests and strict Clippy. Sample-to-flat, merge and frame conversions remain.
This commit is contained in:
Yingwen
+191
-22
@@ -22,7 +22,11 @@ pub mod time;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::any::Any;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use arrow::array::{Array, StringArray};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
|
||||
use crate::error::{InvalidSeriesSnafu, Result};
|
||||
|
||||
@@ -67,10 +71,10 @@ pub struct Label {
|
||||
/// The metric name is only the `__name__` value (see [`Self::name`]); this type is
|
||||
/// the whole label set a PromQL series shows. It is neither a canonical
|
||||
/// matching/grouping identity nor a storage identity (read key / primary key).
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SeriesLabels {
|
||||
name: Option<Arc<str>>,
|
||||
tags: Arc<[Label]>,
|
||||
tags: LabelStorage,
|
||||
}
|
||||
|
||||
impl serde::Serialize for SeriesLabels {
|
||||
@@ -78,7 +82,7 @@ impl serde::Serialize for SeriesLabels {
|
||||
&self,
|
||||
serializer: S,
|
||||
) -> std::result::Result<S::Ok, S::Error> {
|
||||
(&self.name, self.tags.as_ref()).serialize(serializer)
|
||||
(&self.name, self.tags()).serialize(serializer)
|
||||
}
|
||||
}
|
||||
impl<'de> serde::Deserialize<'de> for SeriesLabels {
|
||||
@@ -91,6 +95,129 @@ impl<'de> serde::Deserialize<'de> for SeriesLabels {
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared, bounded series metadata and the owner of its memory reservation.
|
||||
/// Columns map logical label names to UTF-8 arrays; null labels are absent.
|
||||
pub struct SeriesLabelBatch {
|
||||
batch: RecordBatch,
|
||||
columns: Vec<(Arc<str>, usize)>,
|
||||
_owner: Arc<dyn Any + Send + Sync>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for SeriesLabelBatch {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("SeriesLabelBatch")
|
||||
.field("rows", &self.batch.num_rows())
|
||||
.field("columns", &self.columns)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl SeriesLabelBatch {
|
||||
/// Validates mappings once. The owner must account for the retained backing buffers.
|
||||
pub fn try_new(
|
||||
batch: RecordBatch,
|
||||
mut columns: Vec<(Arc<str>, usize)>,
|
||||
owner: Arc<dyn Any + Send + Sync>,
|
||||
) -> Result<Self> {
|
||||
columns.sort_unstable_by(|a, b| a.0.cmp(&b.0));
|
||||
if columns.iter().any(|(name, index)| {
|
||||
name.as_ref() == "__name__"
|
||||
|| batch
|
||||
.columns()
|
||||
.get(*index)
|
||||
.is_none_or(|a| !a.as_any().is::<StringArray>())
|
||||
}) || columns.windows(2).any(|p| p[0].0 == p[1].0)
|
||||
{
|
||||
return invalid_series("invalid batch label mapping");
|
||||
}
|
||||
Ok(Self {
|
||||
batch,
|
||||
columns,
|
||||
_owner: owner,
|
||||
})
|
||||
}
|
||||
|
||||
/// Estimates retained buffers and one row view per batch row, including lazy
|
||||
/// owned-label materialization. Shared buffers are conservatively counted per
|
||||
/// array. Additional independent views of a row need their own reservation.
|
||||
/// This is a reservation estimate, not an allocator-exact accounting contract.
|
||||
pub fn estimated_size(&self) -> usize {
|
||||
let arc_header = 2 * std::mem::size_of::<usize>();
|
||||
let row_views = self
|
||||
.batch
|
||||
.num_rows()
|
||||
.saturating_mul(std::mem::size_of::<BatchLabels>() + arc_header);
|
||||
let mut bytes = self
|
||||
.batch
|
||||
.get_array_memory_size()
|
||||
.saturating_add(std::mem::size_of::<Self>() + arc_header)
|
||||
.saturating_add(
|
||||
self.columns
|
||||
.capacity()
|
||||
.saturating_mul(std::mem::size_of::<(Arc<str>, usize)>()),
|
||||
)
|
||||
.saturating_add(row_views);
|
||||
for (name, column) in &self.columns {
|
||||
bytes = bytes.saturating_add(name.len() + arc_header);
|
||||
for row in 0..self.batch.num_rows() {
|
||||
if let Some(value) = self.value(*column, row) {
|
||||
// Vec-to-Arc conversion may temporarily retain both label arrays.
|
||||
bytes = bytes
|
||||
.saturating_add(2 * std::mem::size_of::<Label>())
|
||||
.saturating_add(value.len())
|
||||
.saturating_add(arc_header + std::mem::align_of::<usize>());
|
||||
}
|
||||
}
|
||||
}
|
||||
bytes.saturating_add(self.batch.num_rows().saturating_mul(arc_header))
|
||||
}
|
||||
|
||||
fn value(&self, column: usize, row: usize) -> Option<&str> {
|
||||
let strings = self
|
||||
.batch
|
||||
.column(column)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()?;
|
||||
(!strings.is_null(row)).then(|| strings.value(row))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
enum LabelStorage {
|
||||
Owned(Arc<[Label]>),
|
||||
Batch(Arc<BatchLabels>),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct BatchLabels {
|
||||
batch: Arc<SeriesLabelBatch>,
|
||||
row: usize,
|
||||
// Existing consumers requesting owned labels share one lazy materialization.
|
||||
materialized: OnceLock<Arc<[Label]>>,
|
||||
}
|
||||
|
||||
impl PartialEq for SeriesLabels {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
if self.name != other.name {
|
||||
return false;
|
||||
}
|
||||
if let (LabelStorage::Batch(a), LabelStorage::Batch(b)) = (&self.tags, &other.tags)
|
||||
&& Arc::ptr_eq(a, b)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
self.tags() == other.tags()
|
||||
}
|
||||
}
|
||||
impl Eq for SeriesLabels {}
|
||||
|
||||
impl std::hash::Hash for SeriesLabels {
|
||||
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
||||
std::hash::Hash::hash(&self.name, state);
|
||||
std::hash::Hash::hash(self.tags(), state);
|
||||
}
|
||||
}
|
||||
|
||||
impl SeriesLabels {
|
||||
/// Creates a metric name, sorting labels and rejecting duplicates or `__name__` tags.
|
||||
pub fn try_new(name: Option<Arc<str>>, mut tags: Vec<Label>) -> Result<Self> {
|
||||
@@ -103,7 +230,7 @@ impl SeriesLabels {
|
||||
}
|
||||
Ok(Self {
|
||||
name,
|
||||
tags: tags.into(),
|
||||
tags: LabelStorage::Owned(tags.into()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -114,18 +241,60 @@ impl SeriesLabels {
|
||||
|
||||
/// Returns ordinary labels sorted by name.
|
||||
pub fn tags(&self) -> &[Label] {
|
||||
&self.tags
|
||||
match &self.tags {
|
||||
LabelStorage::Owned(tags) => tags,
|
||||
LabelStorage::Batch(labels) => labels.materialized.get_or_init(|| {
|
||||
labels
|
||||
.batch
|
||||
.columns
|
||||
.iter()
|
||||
.filter_map(|(name, column)| {
|
||||
labels.batch.value(*column, labels.row).map(|value| Label {
|
||||
name: name.clone(),
|
||||
value: Arc::from(value),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.into()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a shared row view without encoding keys or allocating owned label strings.
|
||||
pub fn from_batch(
|
||||
name: Option<Arc<str>>,
|
||||
batch: Arc<SeriesLabelBatch>,
|
||||
row: usize,
|
||||
) -> Result<Self> {
|
||||
if row >= batch.batch.num_rows() {
|
||||
return invalid_series("series label row out of bounds");
|
||||
}
|
||||
Ok(Self {
|
||||
name,
|
||||
tags: LabelStorage::Batch(Arc::new(BatchLabels {
|
||||
batch,
|
||||
row,
|
||||
materialized: OnceLock::new(),
|
||||
})),
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns a label value, distinguishing absence from an empty value.
|
||||
pub fn get(&self, label: &str) -> Option<&str> {
|
||||
if label == "__name__" {
|
||||
self.name()
|
||||
} else if let LabelStorage::Batch(labels) = &self.tags {
|
||||
let i = labels
|
||||
.batch
|
||||
.columns
|
||||
.binary_search_by(|(name, _)| name.as_ref().cmp(label))
|
||||
.ok()?;
|
||||
labels.batch.value(labels.batch.columns[i].1, labels.row)
|
||||
} else {
|
||||
self.tags
|
||||
self.tags()
|
||||
.binary_search_by(|tag| tag.name.as_ref().cmp(label))
|
||||
.ok()
|
||||
.map(|index| self.tags[index].value.as_ref())
|
||||
.map(|index| self.tags()[index].value.as_ref())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -136,16 +305,16 @@ impl SeriesLabels {
|
||||
return Ok(());
|
||||
}
|
||||
match self
|
||||
.tags
|
||||
.tags()
|
||||
.binary_search_by(|tag| tag.name.as_ref().cmp(label))
|
||||
{
|
||||
Ok(index) => {
|
||||
let mut tags = self.tags.to_vec();
|
||||
let mut tags = self.tags().to_vec();
|
||||
tags[index].value = value;
|
||||
self.tags = tags.into();
|
||||
self.tags = LabelStorage::Owned(tags.into());
|
||||
}
|
||||
Err(index) => {
|
||||
let mut tags = self.tags.to_vec();
|
||||
let mut tags = self.tags().to_vec();
|
||||
tags.insert(
|
||||
index,
|
||||
Label {
|
||||
@@ -153,7 +322,7 @@ impl SeriesLabels {
|
||||
value,
|
||||
},
|
||||
);
|
||||
self.tags = tags.into();
|
||||
self.tags = LabelStorage::Owned(tags.into());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
@@ -164,12 +333,12 @@ impl SeriesLabels {
|
||||
if label == "__name__" {
|
||||
self.name = None;
|
||||
} else if let Ok(index) = self
|
||||
.tags
|
||||
.tags()
|
||||
.binary_search_by(|tag| tag.name.as_ref().cmp(label))
|
||||
{
|
||||
let mut tags = self.tags.to_vec();
|
||||
let mut tags = self.tags().to_vec();
|
||||
tags.remove(index);
|
||||
self.tags = tags.into();
|
||||
self.tags = LabelStorage::Owned(tags.into());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -179,13 +348,13 @@ impl SeriesLabels {
|
||||
self.name = None;
|
||||
}
|
||||
let kept: Vec<_> = self
|
||||
.tags
|
||||
.tags()
|
||||
.iter()
|
||||
.filter(|tag| labels.contains(&tag.name.as_ref()))
|
||||
.cloned()
|
||||
.collect();
|
||||
if kept.len() != self.tags.len() {
|
||||
self.tags = kept.into();
|
||||
if kept.len() != self.tags().len() {
|
||||
self.tags = LabelStorage::Owned(kept.into());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,13 +364,13 @@ impl SeriesLabels {
|
||||
self.name = None;
|
||||
}
|
||||
let kept: Vec<_> = self
|
||||
.tags
|
||||
.tags()
|
||||
.iter()
|
||||
.filter(|tag| !labels.contains(&tag.name.as_ref()))
|
||||
.cloned()
|
||||
.collect();
|
||||
if kept.len() != self.tags.len() {
|
||||
self.tags = kept.into();
|
||||
if kept.len() != self.tags().len() {
|
||||
self.tags = LabelStorage::Owned(kept.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -94,7 +94,9 @@ impl RawTimestamps {
|
||||
/// [`RawSeriesChunk::try_new`]; callers crossing a trust boundary must validate input.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RawSeriesChunk {
|
||||
/// Complete encoded identity of the source series, not a TSID or hash alone.
|
||||
/// Complete identity of the source series, not a TSID or hash alone.
|
||||
/// With explicit labels this may be an opaque domain-qualified identity; only
|
||||
/// a stream declaring a label decode specification requires a label-bearing key.
|
||||
pub primary_key: Buffer,
|
||||
/// `None` denotes an encoded key under bound-label mode; `Some` contains complete
|
||||
/// explicit labels and retains their identity.
|
||||
|
||||
@@ -64,10 +64,10 @@ fn metric_name_edits_retain_shared_labels_and_labels_are_copy_on_write() {
|
||||
)
|
||||
.unwrap();
|
||||
let sibling = metric.clone();
|
||||
assert!(Arc::ptr_eq(&metric.tags, &sibling.tags));
|
||||
assert!(std::ptr::eq(metric.tags(), sibling.tags()));
|
||||
|
||||
metric.set("__name__", Arc::from("memory")).unwrap();
|
||||
assert!(Arc::ptr_eq(&metric.tags, &sibling.tags));
|
||||
assert!(std::ptr::eq(metric.tags(), sibling.tags()));
|
||||
assert_eq!(sibling.name(), Some("cpu"));
|
||||
|
||||
metric.set("host", Arc::from("b")).unwrap();
|
||||
@@ -108,3 +108,56 @@ fn keep_only_and_remove_all_handle_metric_name_and_ordinary_labels() {
|
||||
metric.remove_all(&["zone"]);
|
||||
assert!(metric.tags().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_labels_borrow_values_and_retain_owner() {
|
||||
use super::{LabelStorage, SeriesLabelBatch};
|
||||
use arrow::array::{ArrayRef, StringArray};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
let owner = Arc::new(());
|
||||
let weak = Arc::downgrade(&owner);
|
||||
let batch = RecordBatch::try_from_iter(vec![(
|
||||
"host",
|
||||
Arc::new(StringArray::from(vec![Some(""), None, Some("node")])) as ArrayRef,
|
||||
)])
|
||||
.unwrap();
|
||||
let batch =
|
||||
Arc::new(SeriesLabelBatch::try_new(batch, vec![(Arc::from("host"), 0)], owner).unwrap());
|
||||
let empty = SeriesLabels::from_batch(Some(Arc::from("cpu")), batch.clone(), 0).unwrap();
|
||||
let absent = SeriesLabels::from_batch(None, batch.clone(), 1).unwrap();
|
||||
assert!(SeriesLabels::from_batch(None, batch.clone(), 3).is_err());
|
||||
assert_eq!(empty.get("host"), Some(""));
|
||||
assert_eq!(absent.get("host"), None);
|
||||
let LabelStorage::Batch(view) = &empty.tags else {
|
||||
panic!("expected row view")
|
||||
};
|
||||
assert!(view.materialized.get().is_none());
|
||||
let mut clone = empty.clone();
|
||||
clone.set("host", Arc::from("other")).unwrap();
|
||||
assert_eq!(empty.get("host"), Some(""));
|
||||
assert_eq!(clone.get("host"), Some("other"));
|
||||
assert_eq!(
|
||||
empty,
|
||||
SeriesLabels::try_new(Some(Arc::from("cpu")), vec![label("host", "")]).unwrap()
|
||||
);
|
||||
// Representation changes must preserve keyed collections and persisted labels.
|
||||
let owned = SeriesLabels::try_new(Some(Arc::from("cpu")), vec![label("host", "")]).unwrap();
|
||||
let hash = |labels: &SeriesLabels| {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut state = std::collections::hash_map::DefaultHasher::new();
|
||||
labels.hash(&mut state);
|
||||
state.finish()
|
||||
};
|
||||
assert_eq!(hash(&owned), hash(&empty));
|
||||
let encoded = serde_json::to_string(&empty).unwrap();
|
||||
assert_eq!(encoded, serde_json::to_string(&owned).unwrap());
|
||||
assert_eq!(
|
||||
empty,
|
||||
serde_json::from_str::<SeriesLabels>(&encoded).unwrap()
|
||||
);
|
||||
drop(batch);
|
||||
drop(absent);
|
||||
assert!(weak.upgrade().is_some());
|
||||
drop(empty);
|
||||
assert!(weak.upgrade().is_none());
|
||||
}
|
||||
|
||||
+1
-1
@@ -74,7 +74,7 @@ filtered `RecordBatch` stream.
|
||||
including nullable Float64 fields and preserved sequence/delete rows. Its wire
|
||||
contract is `docs/formats/mito-samples.md`; time/ID bounds may overlap across
|
||||
leaves and parts. Both columnar backends share cancellation-safe output ownership
|
||||
in `series_index/artifact_writer.rs`. Format readers do not apply query visibility; the prepared-source adapter applies visibility and the existing merge/dedup pipeline. Eligible sample merges use identity-only keys and omit tag columns; the raw adapter restores label-bearing keys from candidates. Non-separable tag/sample residuals retain the full flat schema. Manual reconciliation covers all snapshot SSTs regardless of maintenance eligibility, but retains the soft disk limit. Manual builds validate columnar sample and tag schemas before reuse. Columnar reconciliation consolidates snapshot SSTs into UTC-day sample generations; catalogs record all source IDs. Grouped reads require complete daily coverage and all sources still visible, open each artifact once, and fall back for source-sequence-pruned scans. SST replacement invalidates a whole affected generation; catalog publication and reader handles govern retirement.
|
||||
in `series_index/artifact_writer.rs`. Format readers do not apply query visibility; the prepared-source adapter applies visibility and the existing merge/dedup pipeline. Eligible sample merges use identity-only keys and omit tag columns; candidates retain memory-accounted series batch/row labels, and the raw adapter attaches explicit labels with region-qualified identities. Scanner preflight declares that representation so lazy streams and planned frames agree. Non-separable tag/sample residuals retain the full flat schema. Manual reconciliation covers all snapshot SSTs regardless of maintenance eligibility, but retains the soft disk limit. Manual builds validate columnar sample and tag schemas before reuse. Columnar reconciliation consolidates snapshot SSTs into UTC-day sample generations; catalogs record all source IDs. Grouped reads require complete daily coverage and all sources still visible, open each artifact once, and fall back for source-sequence-pruned scans. SST replacement invalidates a whole affected generation; catalog publication and reader handles govern retirement.
|
||||
- **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`.
|
||||
- **WAL/memtable encoding** (`wal/`, `memtable/`): breaks replay if changed incompatibly.
|
||||
|
||||
|
||||
@@ -915,6 +915,16 @@ impl RegionScanner for SeriesScan {
|
||||
.map(|(preparation, _)| preparation.schema))
|
||||
}
|
||||
|
||||
fn raw_series_has_explicit_labels(&self, labels: &RawLabelSource) -> Result<bool, BoxedError> {
|
||||
let Some((mut raw, columnar)) = self.prepare_raw_scan(labels)? else {
|
||||
return Ok(false);
|
||||
};
|
||||
Ok(columnar
|
||||
&& prepared::sample_merge_schema(&self.stream_ctx, &mut raw)
|
||||
.map_err(BoxedError::new)?
|
||||
.is_some())
|
||||
}
|
||||
|
||||
fn scan_series_partition(
|
||||
&self,
|
||||
ctx: &QueryScanContext,
|
||||
@@ -936,6 +946,7 @@ impl RegionScanner for SeriesScan {
|
||||
.build(),
|
||||
));
|
||||
}
|
||||
let explicit_labels = self.raw_series_has_explicit_labels(label_source)?;
|
||||
let future = self.prepare_series(ctx, metrics_set, label_source);
|
||||
let schema = preparation.schema;
|
||||
let labels = label_source.clone();
|
||||
@@ -951,6 +962,7 @@ impl RegionScanner for SeriesScan {
|
||||
schema,
|
||||
labels,
|
||||
self.raw_ordering(prepared_columnar),
|
||||
explicit_labels,
|
||||
Box::pin(stream),
|
||||
)));
|
||||
}
|
||||
|
||||
@@ -33,10 +33,10 @@ use datatypes::arrow::array::{
|
||||
UInt64Array,
|
||||
};
|
||||
use datatypes::arrow::buffer::{Buffer, OffsetBuffer};
|
||||
use datatypes::arrow::compute::{cast, filter_record_batch};
|
||||
use datatypes::arrow::compute::{cast, filter_record_batch, take_record_batch};
|
||||
use datatypes::arrow::datatypes::{Schema, SchemaRef, UInt32Type};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::{SeriesLabelBatch, SeriesLabels, SeriesOrdering};
|
||||
use datatypes::value::{Value, ValueRef};
|
||||
use futures::TryStreamExt;
|
||||
use mito_codec::row_converter::SparsePrimaryKeyCodec;
|
||||
@@ -306,6 +306,61 @@ struct Candidate {
|
||||
key: Buffer,
|
||||
merge_key: Buffer,
|
||||
tags: HashMap<u32, Option<String>>,
|
||||
metric: Option<SeriesLabels>,
|
||||
metadata_row: Option<(Arc<RecordBatch>, usize)>,
|
||||
}
|
||||
|
||||
/// Checks every physical tag before retaining only rows that introduce new identities.
|
||||
fn compact_candidate_batch(
|
||||
batch: &RecordBatch,
|
||||
candidates: &BTreeMap<MetricSeriesId, Candidate>,
|
||||
tags: &[(u32, String)],
|
||||
) -> Result<RecordBatch> {
|
||||
let table = batch
|
||||
.column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
|
||||
.and_then(|a| a.as_any().downcast_ref::<UInt32Array>())
|
||||
.ok_or_else(|| invalid("invalid table ID column"))?;
|
||||
let tsid = batch
|
||||
.column_by_name(DATA_SCHEMA_TSID_COLUMN_NAME)
|
||||
.and_then(|a| a.as_any().downcast_ref::<UInt64Array>())
|
||||
.ok_or_else(|| invalid("invalid TSID column"))?;
|
||||
let mut new_rows = HashMap::new();
|
||||
let mut selected = Vec::new();
|
||||
for row in 0..batch.num_rows() {
|
||||
let id = MetricSeriesId {
|
||||
table_id: table.value(row),
|
||||
tsid: tsid.value(row),
|
||||
};
|
||||
let previous = candidates
|
||||
.get(&id)
|
||||
.and_then(|c| c.metadata_row.as_ref())
|
||||
.map(|(batch, row)| (batch.as_ref(), *row))
|
||||
.or_else(|| new_rows.get(&id).map(|row| (batch, *row)));
|
||||
if let Some((previous, previous_row)) = previous {
|
||||
for (_, name) in tags {
|
||||
let current = batch
|
||||
.column_by_name(name)
|
||||
.and_then(|a| a.as_any().downcast_ref::<StringArray>())
|
||||
.ok_or_else(|| invalid("invalid series tag array"))?;
|
||||
let previous = previous
|
||||
.column_by_name(name)
|
||||
.and_then(|a| a.as_any().downcast_ref::<StringArray>())
|
||||
.ok_or_else(|| invalid("incompatible duplicate label metadata"))?;
|
||||
if previous.is_null(previous_row) != current.is_null(row)
|
||||
|| (!current.is_null(row) && previous.value(previous_row) != current.value(row))
|
||||
{
|
||||
return Err(invalid("conflicting labels for the same table ID and TSID"));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
new_rows.insert(id, row);
|
||||
selected.push(row as u64);
|
||||
}
|
||||
}
|
||||
if selected.len() == batch.num_rows() {
|
||||
return Ok(batch.clone());
|
||||
}
|
||||
take_record_batch(batch, &UInt64Array::from(selected)).context(ComputeArrowSnafu)
|
||||
}
|
||||
|
||||
struct ColumnarPreparedSource {
|
||||
@@ -401,6 +456,13 @@ impl ColumnarPreparedSource {
|
||||
.insert("prepare_series_indexes", indexes.len() as u64);
|
||||
let reservation =
|
||||
MemoryConsumer::new("ColumnarPreparedSource").register(&input.scan_memory_pool);
|
||||
let started = Instant::now();
|
||||
let mut raw = raw;
|
||||
let sample_schema = sample_merge_schema(&context, &mut raw)?;
|
||||
preparation_metrics.add_elapsed("columnar_prepare_schema", started.elapsed());
|
||||
preparation_metrics.bytes.insert("prepare_encoded_keys", 0);
|
||||
preparation_metrics.counts.insert("prepare_batch_series", 0);
|
||||
let direct_labels = sample_schema.is_some();
|
||||
let mut candidates = BTreeMap::<MetricSeriesId, Candidate>::new();
|
||||
for id in indexes {
|
||||
let timer = preparation_metrics.timer("columnar_prepare_series_open_wall");
|
||||
@@ -478,6 +540,12 @@ impl ColumnarPreparedSource {
|
||||
.entry("prepare_selected_rows")
|
||||
.or_default() += batch.num_rows() as u64;
|
||||
let started = Instant::now();
|
||||
let mut duplicate_rows = 0_u64;
|
||||
if direct_labels {
|
||||
let rows = batch.num_rows();
|
||||
batch = compact_candidate_batch(&batch, &candidates, &tags)?;
|
||||
duplicate_rows = (rows - batch.num_rows()) as u64;
|
||||
}
|
||||
let table = batch
|
||||
.column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
|
||||
.ok_or_else(|| invalid("missing table ID column"))?
|
||||
@@ -502,39 +570,80 @@ impl ColumnarPreparedSource {
|
||||
Ok((*id, strings))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let shared_batch = direct_labels.then(|| Arc::new(batch.clone()));
|
||||
let label_batch = if direct_labels {
|
||||
let RawLabelSource::PrimaryKey { logical_tags, .. } = &labels else {
|
||||
return Err(invalid("expected bound labels"));
|
||||
};
|
||||
let columns = logical_tags
|
||||
.iter()
|
||||
.map(|tag| {
|
||||
let name = tags
|
||||
.iter()
|
||||
.find(|(id, _)| *id == tag.column_id)
|
||||
.ok_or_else(|| invalid("missing bound tag"))?;
|
||||
Ok((
|
||||
tag.name.clone(),
|
||||
batch
|
||||
.schema()
|
||||
.index_of(&name.1)
|
||||
.map_err(|e| invalid(e.to_string()))?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let owner = Arc::new(
|
||||
MemoryConsumer::new("ColumnarSeriesLabelBatch")
|
||||
.register(&input.scan_memory_pool),
|
||||
);
|
||||
let label_batch =
|
||||
SeriesLabelBatch::try_new(batch.clone(), columns, owner.clone())
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
owner
|
||||
.try_grow(label_batch.estimated_size())
|
||||
.context(crate::error::ScanMemoryExhaustedSnafu)?;
|
||||
Some(Arc::new(label_batch))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut encoded_key_bytes = 0_u64;
|
||||
let mut tag_values = 0_u64;
|
||||
let mut duplicate_rows = 0_u64;
|
||||
let mut batch_series = 0_u64;
|
||||
for row in 0..batch.num_rows() {
|
||||
let id = MetricSeriesId {
|
||||
table_id: table.value(row),
|
||||
tsid: tsid.value(row),
|
||||
};
|
||||
// TODO(yingwen): Separate raw-series identity from label representation.
|
||||
// The PoC formats assume (table_id, tsid) uniqueness; encode once at
|
||||
// this adapter boundary until RawSeriesChunk no longer needs sparse keys.
|
||||
let mut key = Vec::new();
|
||||
// The flat merger still needs the compact sparse ID prefix.
|
||||
// Explicit-label chunks use a region-qualified identity below.
|
||||
let mut key = if direct_labels {
|
||||
metadata.region_id.as_u64().to_be_bytes().to_vec()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
let identity_prefix_len = key.len();
|
||||
let codec = SparsePrimaryKeyCodec::schemaless();
|
||||
codec
|
||||
.encode_internal(id.table_id, id.tsid, &mut key)
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
let merge_key_len = key.len();
|
||||
codec
|
||||
.encode_to_vec(
|
||||
tag_arrays.iter().map(|(id, strings)| {
|
||||
(
|
||||
*id,
|
||||
if strings.is_null(row) {
|
||||
ValueRef::Null
|
||||
} else {
|
||||
ValueRef::String(strings.value(row))
|
||||
},
|
||||
)
|
||||
}),
|
||||
&mut key,
|
||||
)
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
encoded_key_bytes += key.len() as u64;
|
||||
let merge_key_len = key.len() - identity_prefix_len;
|
||||
if !direct_labels {
|
||||
codec
|
||||
.encode_to_vec(
|
||||
tag_arrays.iter().map(|(id, strings)| {
|
||||
(
|
||||
*id,
|
||||
if strings.is_null(row) {
|
||||
ValueRef::Null
|
||||
} else {
|
||||
ValueRef::String(strings.value(row))
|
||||
},
|
||||
)
|
||||
}),
|
||||
&mut key,
|
||||
)
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
encoded_key_bytes += key.len() as u64;
|
||||
}
|
||||
if let Some(previous) = candidates.get(&id) {
|
||||
duplicate_rows += 1;
|
||||
if previous.key.as_slice() != key.as_slice() {
|
||||
@@ -544,16 +653,19 @@ impl ColumnarPreparedSource {
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// Duplicate identities still undergo a full key comparison before
|
||||
// allocating another owned tag map.
|
||||
tag_values += tag_arrays.len() as u64;
|
||||
let mut values = HashMap::with_capacity(tag_arrays.len());
|
||||
for (column_id, strings) in &tag_arrays {
|
||||
values.insert(
|
||||
*column_id,
|
||||
(!strings.is_null(row)).then(|| strings.value(row).to_owned()),
|
||||
);
|
||||
// Encode borrowed Arrow strings first. Duplicate identities still undergo
|
||||
// the full label-conflict check, without allocating another owned tag map.
|
||||
let mut values = HashMap::new();
|
||||
if !direct_labels {
|
||||
values.reserve(tag_arrays.len());
|
||||
for (column_id, strings) in &tag_arrays {
|
||||
values.insert(
|
||||
*column_id,
|
||||
(!strings.is_null(row)).then(|| strings.value(row).to_owned()),
|
||||
);
|
||||
}
|
||||
}
|
||||
tag_values += values.len() as u64;
|
||||
let bytes = key.len()
|
||||
+ values.values().flatten().map(String::len).sum::<usize>()
|
||||
+ values.capacity() * (std::mem::size_of::<(u32, Option<String>)>() + 16)
|
||||
@@ -562,8 +674,24 @@ impl ColumnarPreparedSource {
|
||||
reservation
|
||||
.try_grow(bytes)
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
let metric = if let Some(label_batch) = &label_batch {
|
||||
batch_series += 1;
|
||||
let RawLabelSource::PrimaryKey { metric_name, .. } = &labels else {
|
||||
return Err(invalid("expected bound labels"));
|
||||
};
|
||||
Some(
|
||||
SeriesLabels::from_batch(
|
||||
Some(metric_name.clone()),
|
||||
label_batch.clone(),
|
||||
row,
|
||||
)
|
||||
.map_err(|e| invalid(e.to_string()))?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let key = Buffer::from(key);
|
||||
let merge_key = key.slice_with_length(0, merge_key_len);
|
||||
let merge_key = key.slice_with_length(identity_prefix_len, merge_key_len);
|
||||
candidates.insert(
|
||||
id,
|
||||
Candidate {
|
||||
@@ -571,6 +699,8 @@ impl ColumnarPreparedSource {
|
||||
key,
|
||||
merge_key,
|
||||
tags: values,
|
||||
metric,
|
||||
metadata_row: shared_batch.as_ref().map(|batch| (batch.clone(), row)),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -586,6 +716,10 @@ impl ColumnarPreparedSource {
|
||||
.counts
|
||||
.entry("prepare_duplicate_rows")
|
||||
.or_default() += duplicate_rows;
|
||||
*preparation_metrics
|
||||
.counts
|
||||
.entry("prepare_batch_series")
|
||||
.or_default() += batch_series;
|
||||
preparation_metrics.add_elapsed("columnar_prepare_candidates", started.elapsed());
|
||||
}
|
||||
}
|
||||
@@ -596,10 +730,6 @@ impl ColumnarPreparedSource {
|
||||
let mut candidates: Vec<_> = candidates.into_values().collect();
|
||||
candidates.sort_unstable_by(|a, b| a.key.as_slice().cmp(b.key.as_slice()));
|
||||
preparation_metrics.add_elapsed("columnar_prepare_sort", started.elapsed());
|
||||
let started = Instant::now();
|
||||
let mut raw = raw;
|
||||
let sample_schema = sample_merge_schema(&context, &mut raw)?;
|
||||
preparation_metrics.add_elapsed("columnar_prepare_schema", started.elapsed());
|
||||
Ok(Self {
|
||||
context,
|
||||
labels,
|
||||
@@ -627,7 +757,7 @@ impl PreparedSeriesSource for ColumnarPreparedSource {
|
||||
.ok_or_else(|| BoxedError::new(invalid("series ordinal out of bounds")))?;
|
||||
Ok(SeriesInfo {
|
||||
primary_key: candidate.key.clone(),
|
||||
metric: None,
|
||||
metric: candidate.metric.clone(),
|
||||
})
|
||||
}
|
||||
fn read_series(
|
||||
@@ -729,13 +859,19 @@ impl PreparedSeriesSource for ColumnarPreparedSource {
|
||||
}
|
||||
let candidates = self.candidates.clone();
|
||||
let reservation = self.reservation.clone();
|
||||
Ok(raw::with_primary_keys(stream, move |key| {
|
||||
Ok(raw::with_series_labels(stream, move |key| {
|
||||
// The resolver can outlive the exhausted sample stream.
|
||||
let _reservation = &reservation;
|
||||
let ordinal = candidates
|
||||
.binary_search_by(|candidate| candidate.merge_key.as_slice().cmp(key.as_slice()))
|
||||
.map_err(|_| BoxedError::new(invalid("unknown merged series identity")))?;
|
||||
Ok(candidates[ordinal].key.clone())
|
||||
Ok((
|
||||
candidates[ordinal].key.clone(),
|
||||
candidates[ordinal]
|
||||
.metric
|
||||
.clone()
|
||||
.ok_or_else(|| BoxedError::new(invalid("identity merge is missing labels")))?,
|
||||
))
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -743,7 +879,7 @@ impl PreparedSeriesSource for ColumnarPreparedSource {
|
||||
/// Omits tags only when every predicate is enforced either on the series index
|
||||
/// or on sample columns. Conjunctions can be separated; mixed disjunctions
|
||||
/// retain the ordinary flat adapter.
|
||||
fn sample_merge_schema(
|
||||
pub(crate) fn sample_merge_schema(
|
||||
context: &StreamContext,
|
||||
raw: &mut RawScanPreparation,
|
||||
) -> Result<Option<SchemaRef>> {
|
||||
|
||||
@@ -121,6 +121,17 @@ async fn write(
|
||||
tsid: u64,
|
||||
timestamp: i64,
|
||||
value: Option<f64>,
|
||||
) {
|
||||
write_with_tags(engine, region, tsid, timestamp, value, &["a", "x"]).await;
|
||||
}
|
||||
|
||||
async fn write_with_tags(
|
||||
engine: &MitoEngine,
|
||||
region: &MitoRegionRef,
|
||||
tsid: u64,
|
||||
timestamp: i64,
|
||||
value: Option<f64>,
|
||||
tags: &[&str],
|
||||
) {
|
||||
let metadata = region.version().metadata.clone();
|
||||
let mut schema = vec![
|
||||
@@ -138,7 +149,7 @@ async fn write(
|
||||
},
|
||||
];
|
||||
let mut values = vec![
|
||||
ValueData::BinaryValue(new_sparse_primary_key(&["a", "x"], &metadata, 10, tsid)),
|
||||
ValueData::BinaryValue(new_sparse_primary_key(tags, &metadata, 10, tsid)),
|
||||
ValueData::TimestampMillisecondValue(timestamp),
|
||||
];
|
||||
let request = if let Some(value) = value {
|
||||
@@ -305,7 +316,22 @@ fn rows(chunks: &[RawSeriesChunk]) -> Vec<(Vec<u8>, i64, f64)> {
|
||||
panic!("expected float");
|
||||
};
|
||||
(
|
||||
chunk.primary_key.to_vec(),
|
||||
{
|
||||
let key = if chunk.metric.is_some() {
|
||||
&chunk.primary_key.as_slice()[8..]
|
||||
} else {
|
||||
chunk.primary_key.as_slice()
|
||||
};
|
||||
let (table, tsid) =
|
||||
mito_codec::row_converter::SparsePrimaryKeyCodec::schemaless()
|
||||
.decode_ids(key)
|
||||
.unwrap();
|
||||
[
|
||||
table.to_be_bytes().as_slice(),
|
||||
tsid.to_be_bytes().as_slice(),
|
||||
]
|
||||
.concat()
|
||||
},
|
||||
chunk.time.timestamp(i).unwrap(),
|
||||
value,
|
||||
)
|
||||
@@ -415,6 +441,13 @@ async fn prepared_sources_share_discovery_and_read_independently() {
|
||||
.any(|(p, _)| p.ends_with(".samples"))
|
||||
);
|
||||
let left = source.read_series(0..2).unwrap();
|
||||
assert!(left.label_decode_spec().is_none());
|
||||
assert!(same_labels(left.label_source(), &binding));
|
||||
for ordinal in 0..source.num_series() {
|
||||
let info = source.series_info(ordinal).unwrap();
|
||||
assert_eq!(info.primary_key, prepared.candidates[ordinal].key);
|
||||
assert_eq!(info.metric, prepared.candidates[ordinal].metric);
|
||||
}
|
||||
let right = source.read_series(0..2).unwrap();
|
||||
assert!(
|
||||
!reads
|
||||
@@ -424,7 +457,18 @@ async fn prepared_sources_share_discovery_and_read_independently() {
|
||||
.any(|(p, _)| p.ends_with(".samples"))
|
||||
);
|
||||
let (left, right) = tokio::join!(left.try_collect::<Vec<_>>(), right.try_collect::<Vec<_>>());
|
||||
let expected = rows(&left.unwrap());
|
||||
let left = left.unwrap();
|
||||
let mut decoder = mito_codec::labels::PrimaryKeyLabelDecoder::from_decode_spec(
|
||||
&common_query::series_frame::label_decode_spec(&legacy_binding).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
for (actual, expected) in left.iter().zip(&oracle) {
|
||||
assert_eq!(
|
||||
actual.metric.as_ref().unwrap(),
|
||||
&decoder.decode(expected.primary_key.as_slice()).unwrap()
|
||||
);
|
||||
}
|
||||
let expected = rows(&left);
|
||||
assert_eq!(rows(&oracle), expected);
|
||||
assert_eq!(expected, rows(&right.unwrap()));
|
||||
assert_eq!(
|
||||
@@ -2130,3 +2174,97 @@ async fn identity_merge_preserves_tag_and_mixed_predicates(#[case] last_non_null
|
||||
}
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overlapping_buckets_retain_only_new_label_rows() {
|
||||
let (_env, engine, region) = fixture().await;
|
||||
const LATER: i64 = 6 * 86_400_000;
|
||||
for timestamp in [3000, LATER] {
|
||||
for id in 1..=32 {
|
||||
let tag = format!("tag-{id}");
|
||||
write_with_tags(&engine, ®ion, id, timestamp, Some(1.0), &["a", &tag]).await;
|
||||
}
|
||||
flush_region(&engine, region.region_id, None).await;
|
||||
}
|
||||
write_with_tags(&engine, ®ion, 33, LATER, Some(1.0), &["a", "tag-33"]).await;
|
||||
flush_region(&engine, region.region_id, None).await;
|
||||
build(&engine, ®ion).await;
|
||||
assert!(region.series_index_version().series_indexes.len() >= 2);
|
||||
let scan = scanner(®ion, region.series_index_store.clone().unwrap());
|
||||
scan.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
&labels(scan.metadata()),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let prepared = scan.prepared.source.get().unwrap().as_ref().unwrap();
|
||||
assert_eq!(34, prepared.candidates.len());
|
||||
let mut retained = HashMap::new();
|
||||
for candidate in prepared.candidates.iter() {
|
||||
let (batch, _) = candidate.metadata_row.as_ref().unwrap();
|
||||
retained.insert(Arc::as_ptr(batch), batch.num_rows());
|
||||
let expected = if candidate.id.tsid == 0 {
|
||||
"x".to_string()
|
||||
} else {
|
||||
format!("tag-{}", candidate.id.tsid)
|
||||
};
|
||||
assert_eq!(
|
||||
candidate.metric.as_ref().unwrap().get("tag_1"),
|
||||
Some(expected.as_str())
|
||||
);
|
||||
}
|
||||
assert_eq!(34, retained.values().sum::<usize>());
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::direct(false)]
|
||||
#[case::mixed_fallback(true)]
|
||||
#[tokio::test]
|
||||
async fn overlapping_buckets_reject_conflicting_physical_tags(#[case] mixed: bool) {
|
||||
let (_env, engine, region) = fixture().await;
|
||||
write_with_tags(
|
||||
&engine,
|
||||
®ion,
|
||||
0,
|
||||
6 * 86_400_000,
|
||||
Some(1.0),
|
||||
&["a", "different"],
|
||||
)
|
||||
.await;
|
||||
flush_region(&engine, region.region_id, None).await;
|
||||
build(&engine, ®ion).await;
|
||||
assert!(region.series_index_version().series_indexes.len() >= 2);
|
||||
let mut input = scan_input(
|
||||
®ion,
|
||||
region.series_index_store.clone().unwrap(),
|
||||
SelectedFormat::Columnar,
|
||||
(0..region.version().metadata.column_metadatas.len()).collect(),
|
||||
);
|
||||
if mixed {
|
||||
input.predicate = PredicateGroup::new(
|
||||
®ion.version().metadata,
|
||||
&[col("tag_0").eq(lit("a")).or(col("field_0").gt(lit(0.0)))],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let scan = SeriesScan::new(input, false);
|
||||
let mut binding = labels(scan.metadata());
|
||||
if let RawLabelSource::PrimaryKey { logical_tags, .. } = &mut binding {
|
||||
// The conflict is in an unbound physical tag, which must still be checked.
|
||||
*logical_tags = logical_tags[..1].to_vec().into();
|
||||
}
|
||||
let error = scan
|
||||
.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
&binding,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.expect("conflicting physical tags must fail preparation");
|
||||
assert!(error.output_msg().contains("conflicting labels"), "{error}");
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
@@ -166,21 +166,26 @@ impl RawSeriesStream for RawStream {
|
||||
}
|
||||
}
|
||||
|
||||
/// Restores the label-bearing key after an identity-only merge has established
|
||||
/// ordering and finality. The resolver must preserve identity and key ordering.
|
||||
pub(crate) fn with_primary_keys(
|
||||
/// Attaches shared labels after the identity-only merge establishes ordering and finality.
|
||||
/// The resolver must preserve identity and key ordering; output keys are opaque identities.
|
||||
pub(crate) fn with_series_labels(
|
||||
stream: SendableRawSeriesStream,
|
||||
resolve: impl Fn(Buffer) -> Result<Buffer, BoxedError> + Send + Sync + 'static,
|
||||
resolve: impl Fn(Buffer) -> Result<(Buffer, datatypes::series::SeriesLabels), BoxedError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> SendableRawSeriesStream {
|
||||
Box::pin(RawStream {
|
||||
schema: stream.schema().clone(),
|
||||
label_source: stream.label_source().clone(),
|
||||
label_decode_spec: stream.label_decode_spec().cloned(),
|
||||
label_decode_spec: None,
|
||||
ordering: stream.ordering(),
|
||||
inner: Box::pin(stream.and_then(move |mut chunk| {
|
||||
let key = resolve(chunk.primary_key.clone());
|
||||
async move {
|
||||
chunk.primary_key = key?;
|
||||
let (key, labels) = key?;
|
||||
chunk.primary_key = key;
|
||||
chunk.metric = Some(labels);
|
||||
Ok(chunk)
|
||||
}
|
||||
})),
|
||||
@@ -504,11 +509,16 @@ pub(crate) fn from_stream(
|
||||
schema: Arc<SeriesSchema>,
|
||||
label_source: RawLabelSource,
|
||||
ordering: SeriesOrdering,
|
||||
explicit_labels: bool,
|
||||
inner: BoxStream<'static, Result<RawSeriesChunk, BoxedError>>,
|
||||
) -> SendableRawSeriesStream {
|
||||
Box::pin(RawStream {
|
||||
schema,
|
||||
label_decode_spec: label_decode_spec(&label_source),
|
||||
label_decode_spec: if explicit_labels {
|
||||
None
|
||||
} else {
|
||||
label_decode_spec(&label_source)
|
||||
},
|
||||
label_source,
|
||||
ordering,
|
||||
inner,
|
||||
|
||||
@@ -2206,7 +2206,7 @@ impl FusedSeriesExec {
|
||||
|
||||
fn is_complete_series_input(plan: &dyn ExecutionPlan) -> bool {
|
||||
if let Some(scan) = plan.downcast_ref::<SeriesScanExec>() {
|
||||
return bound_tags_cover_primary_key(scan.label_source());
|
||||
return bound_tags_cover_primary_key(scan.scan_binding());
|
||||
}
|
||||
if let Some(cooperative) = plan.downcast_ref::<CooperativeExec>() {
|
||||
return cooperative
|
||||
|
||||
@@ -485,7 +485,8 @@ pub struct BoundTag {
|
||||
/// Bound source for materializing labels from raw series chunks.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum RawLabelSource {
|
||||
/// Labels are encoded in the complete primary key and use this scan-bound metadata.
|
||||
/// Scan-bound metadata for labels. A producer may expose encoded keys or
|
||||
/// complete explicit labels; its stream decode specification distinguishes them.
|
||||
PrimaryKey {
|
||||
/// Metric name associated with the bound logical source.
|
||||
metric_name: Arc<str>,
|
||||
@@ -526,10 +527,9 @@ pub const NO_BOUND_LABEL_SOURCE: RawLabelSource = RawLabelSource::Explicit;
|
||||
pub trait RawSeriesStream: Stream<Item = Result<RawSeriesChunk, BoxedError>> + Send {
|
||||
/// Returns the schema shared by all chunks; raw schemas have no evaluation grid.
|
||||
fn schema(&self) -> &Arc<SeriesSchema>;
|
||||
/// Returns scan-bound label metadata. `PrimaryKey` chunks leave `metric` as `None`, meaning
|
||||
/// labels are undecoded, not that the metric name is missing; `Explicit` chunks supply
|
||||
/// `Some` with the complete name/label set, not a partial projection. The metric name itself
|
||||
/// may be absent.
|
||||
/// Returns scan-bound label metadata. The binding does not dictate representation:
|
||||
/// chunks with a decode specification leave `metric` as `None`; otherwise they
|
||||
/// supply the complete explicit label set. The metric name itself may be absent.
|
||||
///
|
||||
/// A stream that was decoded from a raw frame has no scan-bound identity and reports
|
||||
/// [`NO_BOUND_LABEL_SOURCE`], which reads as `Explicit` while its chunks may still carry no
|
||||
@@ -560,7 +560,8 @@ pub type SendableRawSeriesStream = Pin<Box<dyn RawSeriesStream>>;
|
||||
/// Stable identity metadata for one ordinal in a prepared series source.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SeriesInfo {
|
||||
/// Complete encoded primary key, not a TSID or hash alone.
|
||||
/// Complete source identity, not a TSID or hash alone. Explicit-label sources
|
||||
/// may use an opaque domain-qualified identity instead of a label-bearing key.
|
||||
pub primary_key: datatypes::arrow::buffer::Buffer,
|
||||
/// Complete explicit metric name and labels, not a partial label projection; or `None` for
|
||||
/// primary-key label binding (undecoded, not nameless). The metric name itself may be absent.
|
||||
@@ -574,8 +575,9 @@ pub struct SeriesInfo {
|
||||
/// [`PreparedSeriesSource::series_info`], and that identity matches chunks for that ordinal.
|
||||
/// Every reader exposes the same fixed schema, label binding, and ordering before polling,
|
||||
/// including for an empty range, and describes labels independently of this source's lifetime.
|
||||
/// `PrimaryKey` output has `metric: None` (undecoded); `Explicit` output has a complete
|
||||
/// `Some(metric)` matching each chunk. Ordering describes the stream, not ordinal enumeration.
|
||||
/// Key-encoded output has `metric: None`; explicit-label output has a complete
|
||||
/// `Some(metric)` matching each chunk, independently of its scan binding.
|
||||
/// Ordering describes the stream, not ordinal enumeration.
|
||||
///
|
||||
/// Preparation, identity lookup, and reader creation do not load sample payloads or perform
|
||||
/// hidden blocking I/O. The selected set is bound to the scanner's immutable MVCC view,
|
||||
@@ -651,6 +653,13 @@ pub trait RegionScanner: Debug + DisplayAs + Send {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Whether raw chunks carry complete explicit labels under this scan binding.
|
||||
/// This non-consuming declaration must agree with the stream before its first poll.
|
||||
/// Returning true does not replace the immutable source binding or authorization checks.
|
||||
fn raw_series_has_explicit_labels(&self, _labels: &RawLabelSource) -> Result<bool, BoxedError> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// Returns the guarantee for subsequent raw scans without consuming scanner state.
|
||||
///
|
||||
/// Legacy raw sources have no ordering declaration, so `Contiguous` is the conservative
|
||||
|
||||
@@ -622,7 +622,15 @@ impl RegionScanExec {
|
||||
"a scan cannot plan interleaved raw series ordering".into(),
|
||||
));
|
||||
}
|
||||
let arrow_schema = raw_series_frame_schema_with_ordering(&schema, &label_source, ordering)?;
|
||||
let frame_labels = if scanner
|
||||
.raw_series_has_explicit_labels(&label_source)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?
|
||||
{
|
||||
RawLabelSource::Explicit
|
||||
} else {
|
||||
label_source.clone()
|
||||
};
|
||||
let arrow_schema = raw_series_frame_schema_with_ordering(&schema, &frame_labels, ordering)?;
|
||||
let partition_count = scanner.properties().num_partitions().max(1);
|
||||
let properties = Arc::new(PlanProperties::new(
|
||||
EquivalenceProperties::new(arrow_schema.clone()),
|
||||
@@ -633,7 +641,8 @@ impl RegionScanExec {
|
||||
drop(scanner);
|
||||
Ok(Some(Arc::new(SeriesScanExec {
|
||||
scan: self,
|
||||
label_source,
|
||||
label_source: frame_labels,
|
||||
scan_binding: label_source,
|
||||
series_schema: schema,
|
||||
arrow_schema,
|
||||
ordering,
|
||||
@@ -645,8 +654,9 @@ impl RegionScanExec {
|
||||
/// Native physical leaf that packs a prepared region scanner's raw series chunks.
|
||||
/// Do not prepare, reset, or execute concurrently through scanner aliases.
|
||||
pub struct SeriesScanExec {
|
||||
scan: Arc<RegionScanExec>,
|
||||
label_source: RawLabelSource,
|
||||
scan: Arc<RegionScanExec>,
|
||||
scan_binding: RawLabelSource,
|
||||
series_schema: Arc<SeriesSchema>,
|
||||
arrow_schema: ArrowSchemaRef,
|
||||
ordering: SeriesOrdering,
|
||||
@@ -655,7 +665,7 @@ pub struct SeriesScanExec {
|
||||
|
||||
impl std::fmt::Debug for SeriesScanExec {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let source = match &self.label_source {
|
||||
let source = match &self.scan_binding {
|
||||
RawLabelSource::Explicit => "explicit",
|
||||
RawLabelSource::PrimaryKey { .. } => "primary_key",
|
||||
};
|
||||
@@ -671,6 +681,12 @@ impl std::fmt::Debug for SeriesScanExec {
|
||||
}
|
||||
|
||||
impl SeriesScanExec {
|
||||
/// Returns the immutable label binding supplied to the storage scan.
|
||||
pub fn scan_binding(&self) -> &RawLabelSource {
|
||||
&self.scan_binding
|
||||
}
|
||||
|
||||
/// Returns the label representation carried by output frames.
|
||||
pub fn label_source(&self) -> &RawLabelSource {
|
||||
&self.label_source
|
||||
}
|
||||
@@ -721,7 +737,7 @@ fn same_label_source(left: &RawLabelSource, right: &RawLabelSource) -> bool {
|
||||
|
||||
impl DisplayAs for SeriesScanExec {
|
||||
fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
let source = match &self.label_source {
|
||||
let source = match &self.scan_binding {
|
||||
RawLabelSource::Explicit => "explicit",
|
||||
RawLabelSource::PrimaryKey { .. } => "primary_key",
|
||||
};
|
||||
@@ -797,7 +813,7 @@ impl ExecutionPlan for SeriesScanExec {
|
||||
};
|
||||
let scanner = self.scan.scanner.lock().unwrap();
|
||||
let current = scanner
|
||||
.raw_series_schema(&self.label_source)
|
||||
.raw_series_schema(&self.scan_binding)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?
|
||||
.ok_or_else(|| {
|
||||
DataFusionError::Execution("raw series capability changed after planning".into())
|
||||
@@ -805,8 +821,16 @@ impl ExecutionPlan for SeriesScanExec {
|
||||
let current_ordering = scanner
|
||||
.raw_series_ordering()
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let frame_labels = if scanner
|
||||
.raw_series_has_explicit_labels(&self.scan_binding)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?
|
||||
{
|
||||
RawLabelSource::Explicit
|
||||
} else {
|
||||
self.scan_binding.clone()
|
||||
};
|
||||
let current_frame =
|
||||
raw_series_frame_schema_with_ordering(¤t, &self.label_source, current_ordering)?;
|
||||
raw_series_frame_schema_with_ordering(¤t, &frame_labels, current_ordering)?;
|
||||
if current_frame != self.arrow_schema
|
||||
|| scanner.properties().num_partitions().max(1) != partitions
|
||||
{
|
||||
@@ -815,7 +839,7 @@ impl ExecutionPlan for SeriesScanExec {
|
||||
));
|
||||
}
|
||||
let raw = scanner
|
||||
.scan_series_partition(&ctx, &self.scan.metric, partition, &self.label_source)
|
||||
.scan_series_partition(&ctx, &self.scan.metric, partition, &self.scan_binding)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?
|
||||
.ok_or_else(|| {
|
||||
DataFusionError::Execution(
|
||||
@@ -835,13 +859,18 @@ impl ExecutionPlan for SeriesScanExec {
|
||||
.zip(self.series_schema.fields.iter())
|
||||
.any(|(a, b)| a.name != b.name)
|
||||
|| raw.schema().eval_grid.is_some()
|
||||
|| !same_label_source(raw.label_source(), &self.label_source)
|
||||
|| !same_label_source(raw.label_source(), &self.scan_binding)
|
||||
{
|
||||
return Err(DataFusionError::Execution(
|
||||
"raw series stream does not match its planned binding".into(),
|
||||
));
|
||||
}
|
||||
let packed = pack_raw_series(raw)?;
|
||||
if packed.schema() != self.arrow_schema {
|
||||
return Err(DataFusionError::Execution(
|
||||
"raw series label representation changed after planning".into(),
|
||||
));
|
||||
}
|
||||
let stream = RecordBatchStreamAdapter::try_new(packed)
|
||||
.map_err(|e| DataFusionError::External(Box::new(BoxedError::new(e))))?;
|
||||
let stream: SendableRecordBatchStream = Box::pin(stream);
|
||||
|
||||
@@ -58,6 +58,7 @@ struct State {
|
||||
stream_schema: Option<Arc<SeriesSchema>>,
|
||||
stream_source: Option<RawLabelSource>,
|
||||
raw_none: bool,
|
||||
explicit_labels: bool,
|
||||
raw_error: bool,
|
||||
fail_first_poll: bool,
|
||||
error_after_poll: Option<usize>,
|
||||
@@ -201,6 +202,9 @@ impl RegionScanner for TestScanner {
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.raw_schema.clone()))
|
||||
}
|
||||
fn raw_series_has_explicit_labels(&self, _source: &RawLabelSource) -> Result<bool, BoxedError> {
|
||||
Ok(self.state.lock().unwrap().explicit_labels)
|
||||
}
|
||||
fn scan_series_partition(
|
||||
&self,
|
||||
_ctx: &QueryScanContext,
|
||||
@@ -227,7 +231,11 @@ impl RegionScanner for TestScanner {
|
||||
.stream_schema
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.raw_schema.clone().unwrap()),
|
||||
label_decode_spec: common_query::series_frame::label_decode_spec(&source),
|
||||
label_decode_spec: if state.explicit_labels {
|
||||
None
|
||||
} else {
|
||||
common_query::series_frame::label_decode_spec(&source)
|
||||
},
|
||||
source,
|
||||
chunks: self.chunks.clone(),
|
||||
state: self.state.clone(),
|
||||
@@ -1391,3 +1399,35 @@ async fn raw_series_admission_captures_scanner_ordering_before_execution() {
|
||||
assert_eq!(state.drops, 0);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bound_source_can_declare_explicit_labels_and_detect_representation_change() {
|
||||
let schema = series_schema();
|
||||
let (scan, state) = fixture(Some(schema.clone()), vec![chunk(&schema)]);
|
||||
let source = RawLabelSource::PrimaryKey {
|
||||
metric_name: Arc::from("cpu_usage"),
|
||||
physical_metadata: scan.scanner.lock().unwrap().metadata(),
|
||||
logical_tags: Arc::from([]),
|
||||
};
|
||||
state.lock().unwrap().explicit_labels = true;
|
||||
let plan = scan.try_to_series_scan(source).unwrap().unwrap();
|
||||
assert!(matches!(plan.label_source(), RawLabelSource::Explicit));
|
||||
let frames = plan
|
||||
.execute(0, context())
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(frames[0].schema(), plan.schema());
|
||||
let decoded = datatypes::series::frame::unpack_raw_series_frame(
|
||||
&schema,
|
||||
FrameLabelMode::Explicit,
|
||||
None,
|
||||
&frames[0],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(decoded[0].metric, chunk(&schema).metric);
|
||||
state.lock().unwrap().explicit_labels = false;
|
||||
assert!(plan.execute(0, context()).is_err());
|
||||
assert_eq!(state.lock().unwrap().raw_calls, 1);
|
||||
}
|
||||
|
||||
@@ -7460,7 +7460,7 @@ async fn native_promql_rate_and_abs_rate_match_legacy_across_storage_encodings()
|
||||
.expect("fused node reads a prepared series scan");
|
||||
let RawLabelSource::PrimaryKey {
|
||||
physical_metadata, ..
|
||||
} = series_scan.label_source()
|
||||
} = series_scan.scan_binding()
|
||||
else {
|
||||
panic!("fused scan must materialize labels from the primary key");
|
||||
};
|
||||
@@ -7840,7 +7840,7 @@ async fn native_promql_metric_two_aggregate_pair_matches_legacy() {
|
||||
// selector, over the scan bound to that selector's metric.
|
||||
fn scan_metric(plan: &Arc<dyn datafusion::physical_plan::ExecutionPlan>) -> Option<String> {
|
||||
if let Some(scan) = plan.downcast_ref::<SeriesScanExec>() {
|
||||
return match scan.label_source() {
|
||||
return match scan.scan_binding() {
|
||||
store_api::region_engine::RawLabelSource::PrimaryKey {
|
||||
metric_name, ..
|
||||
} => Some(metric_name.to_string()),
|
||||
@@ -8720,10 +8720,14 @@ async fn native_promql_columnar_fused_series_matches_legacy() {
|
||||
metric_name,
|
||||
physical_metadata,
|
||||
logical_tags,
|
||||
} = series_scan.label_source()
|
||||
} = series_scan.scan_binding()
|
||||
else {
|
||||
panic!("the fused scan must materialize labels from the primary key");
|
||||
panic!("the fused scan must bind labels from the primary key");
|
||||
};
|
||||
assert!(matches!(
|
||||
series_scan.label_source(),
|
||||
RawLabelSource::Explicit
|
||||
));
|
||||
assert_eq!(metric_name.as_ref(), TABLE);
|
||||
assert_eq!(
|
||||
logical_tags
|
||||
@@ -8786,6 +8790,19 @@ async fn native_promql_columnar_fused_series_matches_legacy() {
|
||||
let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await;
|
||||
assert_eq!(native_rows, legacy_rows, "same rows on both routes");
|
||||
assert!(observer.fused().is_none(), "the legacy route must not fuse");
|
||||
// An instant selector preserves identities across scan partitions even when frames carry
|
||||
// explicit labels. The binding, rather than the frame encoding, proves completeness.
|
||||
observer.clear();
|
||||
let instant = format!("TQL EVAL (120, 120, '1m') {TABLE}");
|
||||
native_result_rows(native_tql_output(&instance, &instant, true).await).await;
|
||||
let instant_plan = observer.take();
|
||||
let fused = instant_plan
|
||||
.downcast_ref::<promql::extension_plan::FusedSeriesExec>()
|
||||
.unwrap();
|
||||
assert!(fused.has_partition_disjoint_identities());
|
||||
let matched = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}{{host=\"b\"}}[2m])");
|
||||
let matched_rows = native_result_rows(native_tql_output(&instance, &matched, true).await).await;
|
||||
assert_eq!(matched_rows, vec![(120_000, "b".to_string(), 80.0)]);
|
||||
}
|
||||
|
||||
/// No published series index: the columnar `Float64` selector must answer from the raw memtable
|
||||
|
||||
Reference in New Issue
Block a user