From a47a61e8fbd04c64f20db0451e03006e7bbce087 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 2 Nov 2026 14:27:44 +0800 Subject: [PATCH 1/2] refactor(mito2): flatten samples artifacts and reset incompatible index caches (#10041) * refactor(mito2): simplify samples paths and version index cache Signed-off-by: evenyag * fix(mito2): preserve index cache root during reset Signed-off-by: evenyag * ci: refresh compatibility release window Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/formats/mito-samples.md | 10 +- docs/formats/mito-series.md | 7 ++ src/mito2/AGENTS.md | 1 + src/mito2/src/series_index/builder.rs | 7 +- src/mito2/src/series_index/samples.rs | 10 +- src/mito2/src/series_index/samples/reader.rs | 4 +- src/mito2/src/series_index/samples/tests.rs | 11 +- src/mito2/src/worker.rs | 103 ++++++++++++++++++- 8 files changed, 131 insertions(+), 22 deletions(-) diff --git a/docs/formats/mito-samples.md b/docs/formats/mito-samples.md index 72f3fc513ec..b09800df81b 100644 --- a/docs/formats/mito-samples.md +++ b/docs/formats/mito-samples.md @@ -16,9 +16,9 @@ Existing `parquet` and `columnar-series` identifiers keep their meaning. One artifact is derived from one immutable source SST. It owns these objects: ```text -/samples///manifest.json -/samples///part-0.samples -/samples///part-1.samples +/samples//manifest.json +/samples//part-0.samples +/samples//part-1.samples ... ``` @@ -429,3 +429,7 @@ New column encodings or compression methods receive new IDs. Structural changes, changed field meanings, or incompatible limits require a new version. Deeper metadata trees require an explicit reader-support and rollout decision; an unchanged layout alone does not make them readable by version-1 implementations. + +The local cache root contains a layout-version marker. Startup clears the entire +series-index cache when this marker is absent or incompatible, then rebuilds +artifacts through existing maintenance. Source SST identity remains in metadata. diff --git a/docs/formats/mito-series.md b/docs/formats/mito-series.md index a8b3adc627a..5dfe8b46eb5 100644 --- a/docs/formats/mito-series.md +++ b/docs/formats/mito-series.md @@ -392,3 +392,10 @@ stored. A block contains at most 8,388,608 series. The manifest is at most 64 Mi Offsets, lengths, and counts must fit their enclosing areas without arithmetic overflow. Counts must not imply allocations larger than the enclosing decoded payload permits. + +## Local cache compatibility + +The local series-index cache has a root `version` marker independent of this +format version. Before workers start, a missing or incompatible marker causes +the entire derived cache (including samples, ranges, and catalogs) to be reset. +Matching-version restarts preserve the cache. diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index f6b2753b100..a3a04486d65 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -69,6 +69,7 @@ filtered `RecordBatch` stream. - **Series format PoC**: samples are rebuildable local SST sidecars. Optional writer interfaces (`format.rs`) integrate with reconciliation, catalog publication, and deferred purge. Defaults retain the original Parquet writer and disable samples. The opt-in columnar backend (`series_index/columnar.rs`) implements streaming OpenDAL writes and pinned format reads; its wire contract is in `docs/formats/mito-series.md`. Artifact objects record paths and sizes, including the manifest, so reads need no stat. Codecs currently run inline pending scheduling benchmarks. Columnar `SeriesScan` raw reads use scanner-wide preparation in `read/series_scan/prepared.rs`; ordinary record-batch scans retain their SST path. Native reader metadata shares the decoded SST cache reservation; maintenance readers bypass this cache. The paired `experimental_series_index_format` setting selects legacy or columnar artifacts. - **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. +- **Local index cache**: startup checks a root layout version before workers start; missing or incompatible versions reset the whole rebuildable cache, including catalogs and range indexes. Samples use one generation directory; source SST identity stays in metadata. - **Columnar samples**: `series_index/samples.rs` implements opt-in per-SST sidecars, including nullable Float64 fields and preserved sequence/delete rows. Its wire contract is `docs/formats/mito-samples.md`; time/ID bounds may overlap across diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index bd57d362b1d..849c6356011 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -277,11 +277,8 @@ pub(crate) async fn build_samples( )?; let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; let generation = store_api::storage::FileId::random(); - let prefix = format!( - "{}/samples/{}/{generation}/", - region.region_id.as_u64(), - file_id.file_id() - ); + let prefix = + crate::series_index::samples::samples_prefix(region.region_id.as_u64(), generation); let mut writer = backend .writer( version.metadata.clone(), diff --git a/src/mito2/src/series_index/samples.rs b/src/mito2/src/series_index/samples.rs index a5705df07df..8028196513a 100644 --- a/src/mito2/src/series_index/samples.rs +++ b/src/mito2/src/series_index/samples.rs @@ -172,10 +172,7 @@ impl Manifest { } fn prefix(&self) -> String { - format!( - "{}/samples/{}/{}/", - self.region_id, self.source_file_id, self.generation - ) + samples_prefix(self.region_id, self.generation) } fn path(&self) -> String { format!("{}manifest.json", self.prefix()) @@ -233,3 +230,8 @@ fn valid_rows(rows: usize, fields: usize) -> bool { && (1..=MAX_FIELDS).contains(&fields) && rows * (17 + 8 * fields) + fields * rows.div_ceil(8) <= BLOCK_BYTES } + +/// Immutable samples output directory, independent of the source SST identity. +pub(crate) fn samples_prefix(region_id: u64, generation: FileId) -> String { + format!("{region_id}/samples/{generation}/") +} diff --git a/src/mito2/src/series_index/samples/reader.rs b/src/mito2/src/series_index/samples/reader.rs index edea106c498..6bcff2cb814 100644 --- a/src/mito2/src/series_index/samples/reader.rs +++ b/src/mito2/src/series_index/samples/reader.rs @@ -103,8 +103,8 @@ impl SamplesReader { "samples source handle mismatch", )?; let path = format!( - "{region}/samples/{}/{}/manifest.json", - entry.file_id, artifact.generation + "{}manifest.json", + crate::series_index::samples::samples_prefix(region, artifact.generation) ); let mut objects = HashMap::new(); for object in &artifact.objects { diff --git a/src/mito2/src/series_index/samples/tests.rs b/src/mito2/src/series_index/samples/tests.rs index 561569e68e6..8d9ec044d1e 100644 --- a/src/mito2/src/series_index/samples/tests.rs +++ b/src/mito2/src/series_index/samples/tests.rs @@ -126,10 +126,7 @@ async fn build( ) -> SamplesFileHandle { let file_id = FileId::random(); let generation = FileId::random(); - let prefix = format!( - "{}/samples/{file_id}/{generation}/", - meta.region_id.as_u64() - ); + let prefix = format!("{}/samples/{generation}/", meta.region_id.as_u64()); let mut writer = format .writer(meta.clone(), store.clone(), &prefix, file_id, generation) .await @@ -433,7 +430,7 @@ async fn failed_writer_and_abort_remove_outputs() { for drop_writer in [false, true] { let id = FileId::random(); let generation = FileId::random(); - let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + let prefix = format!("{}/samples/{generation}/", m.region_id.as_u64()); let mut writer = format .writer(m.clone(), store.clone(), &prefix, id, generation) .await @@ -934,7 +931,7 @@ async fn wide_schema_byte_bound_and_unsupported_input() { )); let id = FileId::random(); let generation = FileId::random(); - let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + let prefix = format!("{}/samples/{generation}/", m.region_id.as_u64()); assert!( format .writer(wrong, store.clone(), &prefix, id, generation) @@ -1019,7 +1016,7 @@ async fn output_failures_and_cancelled_finish_clean_up() { .unwrap(); let id = FileId::random(); let generation = FileId::random(); - let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + let prefix = format!("{}/samples/{generation}/", m.region_id.as_u64()); let mut writer = format .writer( m.clone(), diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 4bf9d4b244a..de5f41fb4e3 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -526,7 +526,44 @@ async fn series_index_store_from_config( } let root = Path::new(data_home).join("series_index"); - new_fs_cache_store(&root.to_string_lossy()).await.map(Some) + let store = new_fs_cache_store(&root.to_string_lossy()).await?; + initialize_series_index_cache(&store).await?; + Ok(Some(store)) +} + +/// Bump when the rebuildable local index layout becomes incompatible. +const SERIES_INDEX_CACHE_VERSION: &str = "1"; + +/// Resets incompatible caches before any worker or reader can use their catalogs. +async fn initialize_series_index_cache(store: &ObjectStore) -> Result<()> { + let version = match store.read("version").await { + Ok(bytes) => Some(bytes.to_bytes()), + Err(error) if error.kind() == object_store::ErrorKind::NotFound => None, + Err(error) => return Err(error).context(crate::error::OpenDalSnafu), + }; + if version.as_deref() == Some(SERIES_INDEX_CACHE_VERSION.as_bytes()) { + return Ok(()); + } + common_telemetry::info!("Resetting series-index cache: missing or incompatible layout version"); + // Preserve the cache root, which may be a mount point. + let entries = store.list("/").await.context(crate::error::OpenDalSnafu)?; + for entry in entries { + // Listings can include the directory itself. + if entry.path() == "/" { + continue; + } + store + .delete_with(entry.path()) + .recursive(true) + .await + .context(crate::error::OpenDalSnafu)?; + } + store + .write("version", SERIES_INDEX_CACHE_VERSION) + .await + .context(crate::error::OpenDalSnafu)?; + common_telemetry::info!("Series-index cache reset complete"); + Ok(()) } pub async fn write_cache_from_config( @@ -1699,6 +1736,70 @@ mod tests { use super::*; use crate::test_util::TestEnv; + #[tokio::test] + async fn test_series_index_cache_layout_reset() { + let parent = common_test_util::temp_dir::create_temp_dir("series-cache-version"); + let root = parent.path().join("series_index"); + let sentinel = parent.path().join("sst"); + tokio::fs::write(&sentinel, b"source data").await.unwrap(); + let store = crate::access_layer::new_fs_cache_store(&root.to_string_lossy()) + .await + .unwrap(); + // Keep the original directory alive so its inode cannot be reused after deletion. + #[cfg(unix)] + let root_handle = tokio::fs::File::open(&root).await.unwrap(); + for version in [None, Some("0"), Some("invalid")] { + if let Some(version) = version { + store.write("version", version).await.unwrap(); + } + store.write("stale", "old").await.unwrap(); + store + .write("42/samples/old/source/part-0.samples", "old") + .await + .unwrap(); + store + .write("42/samples.json", "stale references") + .await + .unwrap(); + super::initialize_series_index_cache(&store).await.unwrap(); + assert!(!store.exists("stale").await.unwrap()); + assert_eq!( + store.read("version").await.unwrap().to_bytes().as_ref(), + SERIES_INDEX_CACHE_VERSION.as_bytes() + ); + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + + let original = root_handle.metadata().await.unwrap(); + let current = tokio::fs::metadata(&root).await.unwrap(); + assert_eq!( + (original.dev(), original.ino()), + (current.dev(), current.ino()) + ); + } + assert!(!store.exists("42/samples.json").await.unwrap()); + assert!( + !store + .exists("42/samples/old/source/part-0.samples") + .await + .unwrap() + ); + store + .write("42/series/current/manifest.json", "current") + .await + .unwrap(); + super::initialize_series_index_cache(&store).await.unwrap(); + assert!( + store + .exists("42/series/current/manifest.json") + .await + .unwrap() + ); + assert_eq!(tokio::fs::read(&sentinel).await.unwrap(), b"source data"); + } + } + #[test] fn test_region_id_to_index() { let num_workers = 4; From f73a27d779e2fa5e7894c55a03928cd630cfecf5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 2 Nov 2026 15:08:44 +0800 Subject: [PATCH 2/2] feat(promql): own raw series identities at the repartition boundary (#10058) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(promql): own raw series identities at the repartition boundary A raw frame's is_last closes one producer's stream, not the series: the same stored primary key may still arrive from another input. Routing such frames through the evaluated path therefore emits several partial series where a consumer expects one. Teach SeriesRepartitionExec to group raw identities. SeriesRouting selects the grouping key and try_new_with_routing takes an explicit policy, rejecting the combinations a frame's content cannot honour: an evaluated frame may not use RawPrimaryKey, a raw frame may not use CanonicalLabels, and a raw frame must declare explicit labels. Raw routing keys on the complete stored primary key rather than on labels so two pieces of one physical key always meet in one owner, where a disagreement about labels is rejected instead of being split across owners that would each look consistent on its own. A raw identity is published only after every input reaches normal EOF, unlike an evaluated identity, which is complete once each input either contributed it or ended. merge_raw_identity orders the pieces by timestamp, collapses equal duplicates, rejects conflicting ones and emits one closing chunk. The reconstruction's pool charge covers the two parallel row vectors, the timestamps and the worst-case per-row sample slices that alternating contributions produce. This adds the capability only; no node is constructed with raw routing yet. Signed-off-by: discord9 * test(promql): cover raw identity merging and routing source/merge covers the pure merge: interleaved ordering, duplicate collapse versus bit-level conflict, key and label rejection, histogram property comparison and shape mismatch. The router needs its own proof, because the readiness rule and the identity key are what a consumer depends on. A raw child whose partitions are scripted delivers them: one gated input that has already handed over every chunk but whose stream is still open blocks publication until it ends; identical labels with different primary keys stay apart; pieces of one key that disagree on labels fail; and one key with conflicting samples at one timestamp surfaces an error rather than dropping data. The gate holds back only the end of the stream, so the first test distinguishes "every input contributed" from "every input ended". Swapping the raw readiness rule for the evaluated one makes it fail. Signed-off-by: discord9 * docs(promql): correct the cross-partition merge boundary The guide said native execution does not merge one logical series split across partitions. That stopped being true when the repartition owner began finalizing an exact identity's fragments: the merge now happens at that boundary, and the remaining limit is a logical series split across regions. Signed-off-by: discord9 * fix(promql): charge raw chunks their slice and publish complete identities Two problems in the raw path, both about holding memory longer than the data requires. A frame batch decodes into many chunks that are slices of the same Arrow buffers. Charging each chunk the whole allocation charged one batch's timestamps once per chunk: 64 series of 64 rows retained 4.2 MiB for a 32 KiB timestamp buffer, 129 times the real size, and the reservations live until the identity is published. raw_chunk_slice_bytes charges the slice each chunk actually owns, using the same slice-aware accounting the histogram path already used. The same fixture now retains 95 KiB. A raw identity also waited for every input to reach EOF even when every input had already handed over its complete chain for it. read_raw_chains hands a run of chunks over atomically once is_last arrives, and an input cannot contribute the same key twice, so a contribution is final for that input. Readiness is now the same rule for both contents: every input has contributed the identity or ended. The previous test asserted that a fully contributed identity must still wait for an unrelated input's EOF, which pinned the conservative rule as if it were required. It is replaced by two tests, one for each direction: an identity whose inputs all contributed publishes without waiting, and an identity with an input that has contributed nothing waits. Reverting the readiness rule to the old one fails the first and leaves the second passing. Signed-off-by: discord9 * fix(promql): charge histogram indices and steady the raw accounting tests A histogram field is an index array naming rows of a shared table. Charging only the table missed the indices, and a query over few distinct histograms holds mostly indices: 4096 positions naming one row is 16 KiB of indices against a table of a few hundred bytes. Both the histogram and the mixed branch now charge the index array's slice, as the evaluated path already does. The two accounting tests waited by sampling a fixed number of idle ticks. Under load the input task may not have run yet, so the sample reads zero and the test fails for a reason that has nothing to do with what it measures. Withdrawing the index charge reproduces the failure, so the test still discriminates after the wait was made deterministic. Signed-off-by: discord9 * test(promql): gate the producer so the failure cannot drain unclaimed queues The retention this test reads is not guaranteed across `start` and `claim`. `start` spawns the producer before it returns, and a failure drains every queue nobody has claimed yet, releasing the charges those queues held. A producer that reaches its error first therefore leaves the test claiming two empty queues, and "queued output chains retain memory before drain" can never become true. That is the shape the suite now hits on this machine: the baseline commit fails this test in about ten seconds when the suite runs concurrently with another one, while the same test written with the gate passes every time. The producer now waits on a switch before its error, and the test claims both queues, waits for the retained charge, and only then opens the switch. The assertion keeps its meaning because nothing can drain the queues before the test owns them. This is a test fix. The production cleanup it exposed is correct behaviour: a queue no consumer claimed should not keep holding memory. Signed-off-by: discord9 * feat(promql): route raw frames by a match key A vector operation between two raw sources needs the series of one match key in one partition, and the node could not give it that: raw frames were restricted to key routing, so `100 * (1 - a / b)` could not be planned at all. The restriction rested on an inference the code made when it started a read: the routing policy named the frame content, so only one policy could mean raw. That inference stops holding once a raw frame may be routed by a match key, and the content is now read from the child's own declaration instead. Choosing a destination and deciding a fragment's identity are separate jobs. A match key picks the partition by hashing the labels the frame already carries -- raw frames are required to declare them -- while the stored key still decides which identity a fragment belongs to. Both keep one key's fragments in one partition, because one key always decodes to one label set, so merging is unaffected. Canonical labels stay refused for raw frames: they include the metric name, which is not a match key. Signed-off-by: discord9 * fix(promql): report a raw chain that cannot pick a destination Two paths hid a destination they could not compute. A matching-routed chunk without labels built an error and then discarded it through `unwrap_or_else`, routing the chunk by its stored key instead; and `output_labels` answered the stored-key policy with nothing at all, so every chain would have gone to one partition silently. Neither is reachable today -- a raw frame is required to carry explicit labels, and a stored-key-routed frame never asks for a label hash -- but a silent fallback turns a defect into a wrong answer. Both now say what went wrong. Signed-off-by: discord9 * test(promql): cover the histogram properties a duplicate check has to compare Merging two raw pieces treats a sample seen twice at one timestamp as a duplicate only when every property of the two matches. The comparison already covered all of them, but its test changed only the float ones, so dropping `reset_hint`, `start_timestamp` or `custom_values` from the comparison left every test passing: conflicting samples would have been discarded as duplicates with no failure. Each of the three now has a case that differs in exactly that property. Removing any of the three comparisons fails the test. The bucket side is deliberately not covered: a legal sample must satisfy `span_lengths` summing to the bucket count with an in-range cursor, so a sample differing in one bucket field is a different shape rather than a different value. Signed-off-by: discord9 * perf(promql): publish an ordered identity as its inputs pass it, and charge what a merge uses A key-then-time identity could only be published after every input ended, so the node buffered a whole region even though the ordering it was protecting is exactly what lets it publish early. An input emits its keys in order and each chain goes to one owner, so the keys one owner sees from one input are non-decreasing too: a greater key from an input proves that input will not contribute a smaller one. The owner records the greatest key seen per input and publishes the pending identity with the smallest key once every input has contributed it, ended, or been seen past it. Readiness is checked in ascending key order, so the output stays ordered. An identity carried by a single piece needs no reconstruction at all: it is marked final and published as it is, which is what the evaluated path already did. A multi-piece identity charges each allocation the reconstruction actually makes — the two row vectors, the timestamps, and one descriptor per surviving run — instead of reserving a worst case per row per field that a well-ordered identity never approaches. The two tests for the ordered path were written against a fixture where the key they expected to be held was already complete, so they asserted the opposite of the rule. They now use fixtures where the held key really has not been passed, and they collect rows per identity because a frame may split one identity across chunks. Signed-off-by: discord9 * fix(promql): drop the uncharged merge wrapper the charging path replaced Routing a raw identity charges each allocation before it is made, so `merge_raw_identity` had no caller left outside the tests and `cargo clippy -p promql --all-targets -- -D warnings` refused it as dead code. Its body is now the test helper that calls the charging form with a no-op reserve, and the production name is gone. Signed-off-by: discord9 * fix(promql): stop scanning settled identities, and keep merge errors typed Three faults in the raw owner, each small: A key-then-time frame publishes only in key order, so if the smallest pending identity is not ready no larger one can be: an input reaches a later key only by passing the earlier one. The rescan applied that predicate to every pending group anyway, so an input that fell behind could have the owner walk the whole map after each arriving chain it could not yet settle. Ask the first entry only. `reserve` returning `ResourcesExhausted` came back through `invalid(...)`, so a memory-pool rejection reached the caller as a malformed-series error with the wrong status. A reservation failure is not a claim about the data; hand the `DataFusionError` through untouched. The merge reserved its row and selection scratch up front and kept holding both after `rows` and `selected` were dropped, since the reservation outlives them until the packer takes the output. Release the pair once the merge is done. The release amount matches the two reservations, so it can only fail if the reservation was already reduced elsewhere, which is why the assertion checks it but the query still succeeds. Signed-off-by: discord9 * test(promql): assert the decode a chain-atomic merge produces This case asserted that unpacking fails with "evaluated series identity changed before its chain closed". It was written to keep that boundary honest until the region merge became chain-atomic, and its own comment asks for this change once that landed - which it has: unpacking now succeeds. The success it asserts is the same property read the other way. Sixteen hosts, one chain each, each closing exactly once, each reporting 1.0 - the rate its own samples produce, since the fixture's value rises one per second. A merge that started interleaving chains again would cut them, and a cut chain cannot report one correct value per host. The case's evaluation grid holds one step, so it says nothing about how a chain spanning several chunks behaves; that is a property of the grid, not of the merge, and the comment now says so instead of claiming four batches per chain. Signed-off-by: discord9 * test(promql): cover the gaps the review named, and drop two false claims Four independent gaps in the raw repartition tests, plus one comment that asserted something the code does not do. `RawScript` reported `SeriesOrdering::Contiguous` regardless of what its plan declared. `pack_raw_series` derives both the frame metadata and the `SequenceValidator` from that one value, so every `KeyThenTime` test packed `Contiguous` batches and validated them under the `Contiguous` contract — `SequenceValidator`'s "raw series keys are not ordered" check never ran. Carry the declared ordering through `RawScripted` and hand it to the stream, so the batches and the check agree with the plan. `raw_matching_routing_groups_by_match_key` asserted only that both identities shared an output, which holds under primary-key routing too because those two keys collide in a three-output layout. Assert the owner against `output_labels` of the match key instead, so routing by the wrong key fails. The test named for waiting asserted early publication through the watermark. Rename it, and cover the other side deterministically with `gate_chunks`: an identity is withheld while another input has delivered nothing, then publishes once the gate opens. The two lifecycle tests set `gate_chunks: true`, which blocks the *first* chunk, so nothing was ever charged and their wait for a reservation could not be satisfied. Let the chunks through and keep the gate closed at EOF: the chain is delivered, charged, parked, and then dropped. `histogram_duplicates_compare_every_property` claimed the bucket side was unreachable because a legal sample must satisfy layout constraints. Only the bucket *count* is constrained; bucket values are free. Add a case per property — schema, count, zero_count, sum, zero_threshold, either side's bucket value, either side's span offset and span length with equal counts, plus the existing reset_hint, start_timestamp and custom_values — so dropping any comparison from `histograms_equal` fails. `native_prefix_multi_partition_series_chains` said the region merges chain-atomically. It does not: `CoalescePartitionsExec` still interleaves partitions batch by batch. The decode passes because #10023 closes a region's pending raw chunk at package completion, so a chain's value and its close are no longer latched apart. Say that, drop the unsupported "four batches per chain" count, and note that a one-point grid gives one value plus a possible empty close. Signed-off-by: discord9 * perf(promql): merge ordered raw pieces without a row index `merge_raw_identity_charging` rebuilt a `(timestamp, chunk, row)` entry for every row of every piece and sorted the lot, then walked it again to drop duplicate timestamps. Reconstructing one series therefore cost `24 * rows` bytes of transient workspace and an `O(n log n)` sort, even when the pieces already arrived in order. The validation above that point has already refused a piece whose own timestamps are not strictly increasing, so a piece never needs sorting internally. What is left to order is only the relation between pieces. When every adjacent pair satisfies `previous.last() < next.first()` the concatenation is already the merged sequence, and no timestamp can repeat across pieces, so neither the index nor the duplicate walk is needed. Keep the general path unchanged for the interleaved and touching cases - touching pieces share a timestamp, so they still need the duplicate check - and shrink only what the path actually reserved. Signed-off-by: discord9 * test(promql): vary histogram spans one dimension at a time `histogram_duplicates_compare_every_property` varied `span_offsets` and `span_lengths` together in the two length cases: the baseline was `offsets: [0], lengths: [2]` and the cases used `offsets: [0, 0], lengths: [1, 1]`. The offsets comparison rejected them first, so removing the lengths comparison left the test green and those two cases pinned nothing. The baseline is now two spans in both sides (`offsets: [0, 0]`, `lengths: [1, 1]`). The offset cases change only the cursors (`[1, 0]`) and the length cases only the lengths (`[2, 0]`), keeping the cursor list, the span count and the bucket count fixed, with a zero-length span as the only difference. Ablating any one of the four comparisons - positive and negative offsets, positive and negative lengths - now fails the case that covers it. `single_chunk_uses_ordered_fast_path` is removed. It passed with the fast path forced off, so it did not pin path selection, and the owner never reaches the merge with one chunk: `finish_raw_group` returns the chunk after setting `is_last`. The fast path stays pinned by `ordered_disjoint_pieces_concatenate_without_row_workspace`, which asserts the exact reservation and zero shrink for disjoint chunks, and by `touching_ordered_chunks_use_general_duplicate_path` on the other side. Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/promql/AGENTS.md | 9 +- src/promql/src/extension_plan.rs | 2 +- .../src/extension_plan/series_repartition.rs | 808 ++++++++++-- .../series_repartition/tests.rs | 1133 ++++++++++++++++- src/promql/src/source.rs | 3 + src/promql/src/source/merge.rs | 252 ++++ src/promql/src/source/merge/tests.rs | 702 ++++++++++ src/promql/src/vector_match.rs | 4 +- tests-integration/tests/native_prefix.rs | 60 +- 9 files changed, 2791 insertions(+), 182 deletions(-) create mode 100644 src/promql/src/source/merge.rs create mode 100644 src/promql/src/source/merge/tests.rs diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index 6af954ee08e..437c58b1315 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -23,10 +23,11 @@ functions, vector matching and extension operators. Statement planning lives in Native PromQL is an experimental, standalone evaluation route alongside the existing DataFusion PromQL path. It operates on one storage region, which may -supply multiple scan partitions. It can evaluate each independent input series -in those partitions, but does not merge one logical series split across regions -or partitions; cross-region `SeriesMerge`/`Exchange` and global final-label -semantics are not implemented. Do not describe partition-local work as a global +supply multiple scan partitions. Where a boundary must answer one logical +series once, `SeriesRepartitionExec` merges that identity's fragments across +scan partitions and recomputes a single close for it. Cross-region +`SeriesMerge`/`Exchange` are still not implemented, so a logical series split +across regions is not merged. Do not describe partition-local work as a global series merge. Vector operations use the existing vector-matching semantics; do not describe diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index ba4cf65e0d5..c8c281a22ff 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -57,7 +57,7 @@ pub use range_manipulate::{RangeManipulate, RangeManipulateExec, RangeManipulate pub use raw_series_source::{RawSeriesSource, RawSeriesSourceExec}; pub use scalar_calculate::ScalarCalculate; pub use series_divide::{SeriesDivide, SeriesDivideExec, SeriesDivideStream}; -pub use series_repartition::SeriesRepartitionExec; +pub use series_repartition::{SeriesRepartitionExec, SeriesRouting}; pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctOnStream}; pub type Millisecond = ::Native; diff --git a/src/promql/src/extension_plan/series_repartition.rs b/src/promql/src/extension_plan/series_repartition.rs index af53ef796d9..fc3c89d68cb 100644 --- a/src/promql/src/extension_plan/series_repartition.rs +++ b/src/promql/src/extension_plan/series_repartition.rs @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Routes evaluated series chains by canonical label set across DataFusion input partitions. +//! Routes evaluated and raw series across DataFusion input partitions. //! -//! This node routes chains by canonical label-set hash, then merges exact metric-name/tag -//! identities at each owner after their close/EOF barriers. This is owner-local grouping, not a claim -//! of global canonical completeness. It operates on DataFusion partitions only; arbitrary `MatchKey` -//! and `GroupKey` routing is unsupported. +//! Canonical routing hashes complete labels; matching routing hashes the evaluator's own PromQL +//! matching key. Routing groups are not merge identities: evaluated series merge only by exact name +//! and tags. Raw ownership is the complete primary key plus exact label agreement; each input's +//! complete chain or EOF establishes its contribution to a raw identity. //! //! Chains are transported intact to owners with their input-partition ID, followed by one charged //! normal-EOF message per live output from each completed input. Owners group exact identities and @@ -121,7 +121,8 @@ use std::task::{Context, Poll}; use common_error::ext::BoxedError; use common_query::series_frame::{ - SendableEvaluatedSeriesStream, pack_evaluated_series, unpack_evaluated_series, + SendableEvaluatedSeriesStream, pack_evaluated_series, pack_raw_series, unpack_evaluated_series, + unpack_raw_series, }; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -135,18 +136,64 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; -use datatypes::series::SeriesLabels; -use datatypes::series::frame::{FrameContent, SeriesFrameDeclaration}; +use datatypes::series::frame::{ + FrameContent, FrameLabelMode, LabelDecodeSpec, SeriesFrameDeclaration, +}; +use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::{Series, SeriesSchema}; +use datatypes::series::{SeriesLabels, SeriesOrdering}; use futures::{Stream, StreamExt}; +use store_api::region_engine::{NO_BOUND_LABEL_SOURCE, RawSeriesStream, SendableRawSeriesStream}; use tokio::sync::mpsc; use twox_hash::XxHash64; use crate::grid_series::{label_set, merge_same_labelset}; -use crate::native::execute::merge_workspace_bytes; -use crate::native::execute::operand_chunk_bytes; +use crate::native::execute::{merge_workspace_bytes, operand_chunk_bytes}; +use crate::source::merge_raw_identity_charging; use datatypes::series::samples::DenseSamplesRef; +/// Selects the key used to route a series to an output partition. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SeriesRouting { + /// Hash the canonical label set. + CanonicalLabels, + /// Hash the exact PromQL matching key, so both sides of a vector operation meet. + Matching(crate::vector_match::MatchingLabels), + /// Route raw frames by their complete stored primary key. + RawPrimaryKey, +} + +struct EmptyRawStream { + schema: Arc, + ordering: SeriesOrdering, +} + +impl Stream for EmptyRawStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(None) + } +} + +impl RawSeriesStream for EmptyRawStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &NO_BOUND_LABEL_SOURCE + } + + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + None + } + + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + /// The seed behind every routing decision. /// /// Fixed rather than drawn per execution, so which output a label set belongs on is a function of @@ -165,6 +212,9 @@ pub struct SeriesRepartitionExec { series_schema: Arc, /// The one hash configuration of this node's execution. routing: Routing, + routing_policy: SeriesRouting, + content: FrameContent, + ordering: SeriesOrdering, /// Started by the first `execute` call, since the output partitions share one read of the child. shared: Mutex>>, properties: Arc, @@ -184,6 +234,25 @@ impl SeriesRepartitionExec { /// and a frame whose declaration does not survive the round trip would be published under a /// schema its batches do not match. pub fn try_new(input: Arc, outputs: usize) -> DataFusionResult { + let declaration = + SeriesFrameDeclaration::parse(input.schema().as_ref()).map_err(|error| { + DataFusionError::Plan(format!( + "SeriesRepartition input is not a series frame: {error}" + )) + })?; + let policy = match declaration.content() { + FrameContent::Evaluated => SeriesRouting::CanonicalLabels, + FrameContent::Raw => SeriesRouting::RawPrimaryKey, + }; + Self::try_new_with_routing(input, outputs, policy) + } + + /// Creates the node with an explicit content-compatible routing policy. + pub fn try_new_with_routing( + input: Arc, + outputs: usize, + routing_policy: SeriesRouting, + ) -> DataFusionResult { if outputs == 0 { return Err(DataFusionError::Plan( "SeriesRepartition requires at least one output partition".to_string(), @@ -200,17 +269,45 @@ impl SeriesRepartitionExec { "SeriesRepartition input is not a series frame: {error}" )) })?; - if declaration.content() != FrameContent::Evaluated { + let content = declaration.content(); + match (content, &routing_policy) { + (FrameContent::Evaluated, SeriesRouting::RawPrimaryKey) => { + return Err(DataFusionError::Plan( + "SeriesRepartition evaluated frames cannot use RawPrimaryKey routing".into(), + )); + } + (FrameContent::Raw, SeriesRouting::CanonicalLabels) => { + return Err(DataFusionError::Plan( + "SeriesRepartition raw frames cannot use CanonicalLabels routing".into(), + )); + } + _ => {} + } + if content == FrameContent::Raw && declaration.labels() != FrameLabelMode::Explicit { return Err(DataFusionError::Plan( - "SeriesRepartition routes evaluated series fragments, not raw samples".to_string(), + "SeriesRepartition raw frames require explicit labels".into(), )); } let series_schema = Arc::new(declaration.series().clone()); - let declared = pack_evaluated_series( - Box::pin(futures::stream::empty()), - Arc::clone(&series_schema), - )? - .schema(); + let ordering = if content == FrameContent::Raw { + common_query::series_frame::series_frame_ordering(input.schema().as_ref())? + } else { + SeriesOrdering::Contiguous + }; + let declared = match content { + FrameContent::Evaluated => pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + )? + .schema(), + FrameContent::Raw => { + let source = EmptyRawStream { + schema: Arc::clone(&series_schema), + ordering, + }; + pack_raw_series(Box::pin(source))?.schema() + } + }; if declared != input.schema() { return Err(DataFusionError::Plan( "SeriesRepartition input frame carries a schema this node cannot re-declare" @@ -231,14 +328,22 @@ impl SeriesRepartitionExec { Ok(Self { input, series_schema, - routing: Routing::new(outputs), + routing: Routing::with_policy(outputs, routing_policy.clone()), + routing_policy, + content, + ordering, shared: Mutex::new(None), properties, }) } + /// The routing policy used by this node. + pub fn routing_policy(&self) -> &SeriesRouting { + &self.routing_policy + } + /// How many output partitions this node has. - fn outputs(&self) -> usize { + pub fn outputs(&self) -> usize { self.properties.partitioning.partition_count() } } @@ -313,9 +418,10 @@ impl ExecutionPlan for SeriesRepartitionExec { } // Rebuilt rather than patched, so a child whose schema changed is refused here instead of // routed under a frame declaration the new child does not make. - Ok(Arc::new(Self::try_new( + Ok(Arc::new(Self::try_new_with_routing( children.swap_remove(0), self.outputs(), + self.routing_policy.clone(), )?)) } @@ -340,7 +446,7 @@ impl ExecutionPlan for SeriesRepartitionExec { // emitted twice or read by a scan that was never meant to run twice. *state = Some(SharedExecution::start( &self.input, - self.routing, + self.routing.clone(), Arc::clone(&self.series_schema), &context, )?); @@ -359,9 +465,25 @@ impl ExecutionPlan for SeriesRepartitionExec { Arc::clone(&self.series_schema), &context.runtime_env().memory_pool, self.input.properties().partitioning.partition_count(), + self.ordering, ); let handoff = Arc::clone(&inner.handoff); - let output = pack_evaluated_series(Box::pin(inner), Arc::clone(&self.series_schema))?; + let output: SendableRecordBatchStream = match self.content { + FrameContent::Evaluated => pack_evaluated_series( + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(inner)), + }), + Arc::clone(&self.series_schema), + )?, + FrameContent::Raw => { + let raw = RawOutputStream { + inner: Some(Box::pin(inner)), + schema: Arc::clone(&self.series_schema), + ordering: self.ordering, + }; + pack_raw_series(Box::pin(raw))? + } + }; let schema = output.schema(); let output = Box::pin(PackedHandoffStream { inner: Some(output), @@ -433,6 +555,11 @@ impl SharedExecution { series_schema: Arc, context: &Arc, ) -> DataFusionResult> { + // What the child carries is what it declares, not what the policy suggests: a raw frame + // can be routed by a match key too, so the policy no longer names the content. + let content = SeriesFrameDeclaration::parse(input.schema().as_ref()) + .map_err(|error| DataFusionError::Plan(error.to_string()))? + .content(); let partitions = input.properties().partitioning.partition_count(); let mut streams = Vec::with_capacity(partitions); for child_partition in 0..partitions { @@ -440,10 +567,15 @@ impl SharedExecution { // The expected schema is a check and never an input: the decoder reads the frame's own // declaration, and a partition declaring another frame is refused here rather than // routed under labels it does not carry. - streams.push(unpack_evaluated_series( - frames, - Some(Arc::clone(&series_schema)), - )?); + streams.push(match content { + FrameContent::Evaluated => RoutedInput::Evaluated(unpack_evaluated_series( + frames, + Some(Arc::clone(&series_schema)), + )?), + FrameContent::Raw => { + RoutedInput::Raw(unpack_raw_series(frames, Some(Arc::clone(&series_schema)))?) + } + }); } let mut senders = Vec::with_capacity(routing.outputs); let mut receivers = Vec::with_capacity(routing.outputs); @@ -487,7 +619,7 @@ impl SharedExecution { let task = common_runtime::spawn_query(route_partition( stream, partition, - routing, + routing.clone(), senders.clone(), // What this partition reads ahead is charged to its own consumer, so the read-ahead // of one partition is attributable to the task that did it. @@ -839,7 +971,7 @@ impl Inputs { /// queued chains remain charged to the pool until consumed or dropped. Each task records its outcome /// before returning and dropping its senders, so output streams observe the final read outcome. async fn route_partition( - mut stream: SendableEvaluatedSeriesStream, + mut stream: RoutedInput, partition: usize, routing: Routing, senders: Vec>, @@ -847,7 +979,15 @@ async fn route_partition( inputs: Arc, mut report: ReportIfUnreported, ) { - let stop = read_chains(&mut stream, partition, &routing, &senders, &read_ahead).await; + // The variant is the content: decoding the input's declared frame is what produced it. + let stop = match &mut stream { + RoutedInput::Evaluated(stream) => { + read_evaluated_chains(stream, partition, &routing, &senders, &read_ahead).await + } + RoutedInput::Raw(stream) => { + read_raw_chains(stream, partition, &routing, &senders, &read_ahead).await + } + }; if matches!(&stop, Stop::Ended) { for sender in &senders { if sender.is_closed() { @@ -963,7 +1103,7 @@ pub(crate) fn histogram_table_slice_bytes( /// to the pool. Output queues are unbounded, so routing does not wait for output polling. Each queued /// chain retains its pool charge until the owner consumes it or the queue is dropped. This function /// otherwise waits only on its own child. -async fn read_chains( +async fn read_evaluated_chains( stream: &mut SendableEvaluatedSeriesStream, partition: usize, routing: &Routing, @@ -973,7 +1113,7 @@ async fn read_chains( loop { // One whole chain, assembled here: the chunk that closes the series ends it, and so does a // stream that ends between two chains. - let mut chain: Chain = Vec::new(); + let mut chain: Vec<(Series, bool, MemoryReservation)> = Vec::new(); let mut chain_output: Option = None; let mut discarding = false; let mut ended = false; @@ -1000,7 +1140,7 @@ async fn read_chains( let output = match chain_output { Some(output) => output, None => { - let output = routing.output(&series.metric); + let output = routing.output_labels(&series.metric); chain_output = Some(output); output } @@ -1095,7 +1235,7 @@ async fn read_chains( senders[output] .send(Queued::Chain { input: partition, - chain, + chain: RoutedChain::Evaluated(chain), _metadata: metadata, }) .is_err() @@ -1106,6 +1246,112 @@ async fn read_chains( } } +async fn read_raw_chains( + stream: &mut SendableRawSeriesStream, + partition: usize, + routing: &Routing, + senders: &[mpsc::UnboundedSender], + read_ahead: &MemoryReservation, +) -> Stop { + loop { + let mut chain: Vec<(RawSeriesChunk, MemoryReservation)> = Vec::new(); + let mut destination: Option = None; + let mut discarding = false; + loop { + tokio::task::coop::consume_budget().await; + let next = match (destination, discarding) { + (Some(output), false) if !chain.is_empty() => tokio::select! { + biased; + _ = senders[output].closed() => { chain.clear(); discarding = true; continue; } + item = stream.next() => item, + }, + _ => stream.next().await, + }; + match next { + Some(Ok(chunk)) => { + let output = match destination { + Some(output) => output, + None => { + let output = match &routing.policy { + SeriesRouting::Matching(_) => { + let Some(labels) = chunk.metric.as_ref() else { + return Stop::Failed(Arc::new(DataFusionError::Execution( + "a matching-routed raw chain carries no labels".into(), + ))); + }; + routing.output_labels(labels) + } + _ => routing.output_key(chunk.primary_key.as_slice()), + }; + *destination.insert(output) + } + }; + let final_chunk = chunk.is_last; + if senders[output].is_closed() { + chain.clear(); + discarding = true; + if final_chunk { + break; + } + continue; + } + if discarding { + if final_chunk { + break; + } + continue; + } + let charge = read_ahead.new_empty(); + let bytes = match raw_chunk_slice_bytes(&chunk) { + Ok(bytes) => bytes, + Err(error) => return Stop::Failed(Arc::new(error)), + }; + if let Err(error) = charge.try_grow(bytes) { + return Stop::Failed(Arc::new(error)); + } + chain.push((chunk, charge)); + if final_chunk { + break; + } + } + Some(Err(error)) => { + return Stop::Failed(Arc::new(DataFusionError::External(Box::new(error)))); + } + None => { + return if destination.is_none() { + Stop::Ended + } else { + Stop::Failed(Arc::new(DataFusionError::Execution( + "raw series ended without is_last".into(), + ))) + }; + } + } + } + let Some(output) = destination else { + return Stop::Failed(Arc::new(DataFusionError::Internal( + "raw chain has no destination".into(), + ))); + }; + if !discarding { + let metadata = read_ahead.new_empty(); + if let Err(error) = metadata.try_grow(std::mem::size_of::()) { + return Stop::Failed(Arc::new(error)); + } + if senders[output] + .send(Queued::Chain { + input: partition, + chain: RoutedChain::Raw(chain), + _metadata: metadata, + }) + .is_err() + { + continue; + } + } + } +} + fn exact_identity(series: &Series) -> ExactIdentity { ( series.metric.name().map(Arc::::from), @@ -1225,13 +1471,21 @@ enum Ending { /// /// The transport unit is a chain so queued source contributions remain intact. Owners collect all /// chains before locally merging exact identities and producing their final output. -type Chain = Vec<(Series, bool, MemoryReservation)>; +enum RoutedInput { + Evaluated(SendableEvaluatedSeriesStream), + Raw(SendableRawSeriesStream), +} + +enum RoutedChain { + Evaluated(Vec<(Series, bool, MemoryReservation)>), + Raw(Vec<(RawSeriesChunk, MemoryReservation)>), +} /// What one output queue carries: a whole chain and its memory charge. enum Queued { Chain { input: usize, - chain: Chain, + chain: RoutedChain, _metadata: MemoryReservation, }, Eof { @@ -1242,14 +1496,199 @@ enum Queued { type ExactIdentity = (Option>, Vec<(Arc, Arc)>); +/// Turns an owner's evaluated publications into the evaluated stream the packer consumes. +struct EvaluatedOutputStream { + inner: Option>>, +} + +impl Stream for EvaluatedOutputStream { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Some(inner) = this.inner.as_mut() else { + return Poll::Ready(None); + }; + match inner.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => { + this.inner = None; + Poll::Ready(None) + } + Poll::Ready(Some(Err(error))) => { + this.inner = None; + Poll::Ready(Some(Err(error))) + } + // A merged identity is final: the owner publishes it once, already closed. + Poll::Ready(Some(Ok(ReadyOutput::Evaluated(series)))) => { + Poll::Ready(Some(Ok((series, true)))) + } + Poll::Ready(Some(Ok(ReadyOutput::Raw(_)))) => { + this.inner = None; + Poll::Ready(Some(Err(boxed(DataFusionError::Internal( + "an evaluated repartition output carried a raw series".into(), + ))))) + } + } + } +} + +/// Turns an owner's raw publications into the raw frame stream the packer consumes. +struct RawOutputStream { + inner: Option>>, + schema: Arc, + ordering: SeriesOrdering, +} + +impl Stream for RawOutputStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Some(inner) = this.inner.as_mut() else { + return Poll::Ready(None); + }; + match inner.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => { + this.inner = None; + Poll::Ready(None) + } + Poll::Ready(Some(Err(error))) => { + this.inner = None; + Poll::Ready(Some(Err(error))) + } + Poll::Ready(Some(Ok(ReadyOutput::Raw(chunk)))) => Poll::Ready(Some(Ok(chunk))), + Poll::Ready(Some(Ok(ReadyOutput::Evaluated(_)))) => { + this.inner = None; + Poll::Ready(Some(Err(boxed(DataFusionError::Internal( + "a raw repartition output carried an evaluated series".into(), + ))))) + } + } + } +} + +impl RawSeriesStream for RawOutputStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &store_api::region_engine::NO_BOUND_LABEL_SOURCE + } + + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + None + } + + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +/// What one owner groups by: an evaluated exact identity, or a complete raw primary key. +/// +/// The two are separate variants rather than one tuple, so a raw key can never be confused with an +/// evaluated identity that happens to look similar. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +enum OwnerKey { + Evaluated(ExactIdentity), + Raw(Vec), +} + +/// The raw identity of one chunk: its complete stored primary key and nothing else. +/// +/// Labels are deliberately not part of the key. Two pieces of one physical key must always meet in +/// the same owner so a label disagreement can be rejected there, instead of being split across +/// owners that would each look consistent on its own. +fn raw_chunk_slice_bytes(chunk: &RawSeriesChunk) -> DataFusionResult { + use datafusion::arrow::array::Array; + + let metric_bytes = chunk.metric.as_ref().map_or(0, |metric| { + metric + .name() + .map_or(0, |name| name.len()) + .saturating_add(metric.tags().iter().fold(0usize, |bytes, tag| { + bytes + .saturating_add(tag.name.len()) + .saturating_add(tag.value.len()) + })) + }); + // A histogram field is an index array naming rows of a shared table; both are held per sample. + let mut fields_bytes = 0usize; + for samples in chunk.fields.iter() { + let field_bytes = match samples.dense() { + Some(DenseSamplesRef::Float(values)) => values.to_data().get_slice_memory_size()?, + Some(DenseSamplesRef::Histogram(histograms)) => histograms + .indices() + .to_data() + .get_slice_memory_size()? + .saturating_add(histogram_table_slice_bytes(histograms.batch())?), + Some(DenseSamplesRef::Mixed(mixed)) => mixed + .floats() + .to_data() + .get_slice_memory_size()? + .saturating_add( + mixed + .histograms() + .indices() + .to_data() + .get_slice_memory_size()?, + ) + .saturating_add(histogram_table_slice_bytes(mixed.histograms().batch())?), + None => 0, + }; + fields_bytes = fields_bytes.saturating_add(field_bytes); + } + Ok(std::mem::size_of::() + .saturating_add(chunk.primary_key.len()) + .saturating_add(metric_bytes) + .saturating_add( + chunk + .time + .as_buffer() + .len() + .saturating_mul(std::mem::size_of::()), + ) + .saturating_add(fields_bytes)) +} + +fn raw_identity(chunk: &RawSeriesChunk) -> OwnerKey { + OwnerKey::Raw(chunk.primary_key.as_slice().to_vec()) +} + struct PendingGroup { fragments: Vec, + raw_chunks: Vec, charges: Vec, received: Vec, metadata: MemoryReservation, } -fn group_ready(group: &PendingGroup, eof: &[bool]) -> bool { +/// Whether this owner may publish `group` now. +/// +/// A key-then-time raw frame promises its keys arrive in order, and this owner publishes in key +/// order only when it publishes them together: an identity that is complete while a smaller one is +/// still waiting would go out first and break that promise. Watermarks let it publish once every +/// input has advanced past the identity, without waiting for every input to end. +fn group_ready( + group: &PendingGroup, + eof: &[bool], + last_key_from: &[Option>], + ordering: SeriesOrdering, + key: &OwnerKey, +) -> bool { + if ordering == SeriesOrdering::KeyThenTime { + let OwnerKey::Raw(key) = key else { + return false; + }; + return group.received.iter().zip(eof).zip(last_key_from).all( + |((received, eof), last_key)| { + *received || *eof || last_key.as_ref().is_some_and(|last| last > key) + }, + ); + } group .received .iter() @@ -1257,8 +1696,13 @@ fn group_ready(group: &PendingGroup, eof: &[bool]) -> bool { .all(|(received, eof)| *received || *eof) } +enum ReadyOutput { + Evaluated(Series), + Raw(RawSeriesChunk), +} + struct ReadySeries { - series: Series, + series: ReadyOutput, charges: ReadyCharges, } @@ -1274,8 +1718,11 @@ struct RepartitionStream { lease: OutputLease, schema: Arc, pool: Arc, - groups: BTreeMap, + /// What the frame this stream emits declares about the order of its keys. + ordering: SeriesOrdering, + groups: BTreeMap, eof: Vec, + last_key_from: Vec>>, eof_metadata: Option, inputs: usize, output: VecDeque, @@ -1288,7 +1735,7 @@ struct RepartitionStream { /// every identity it releases. rescan: bool, /// The last identity a pass examined, exclusive. Retains at most one identity descriptor. - resume_after: Option, + resume_after: Option, reported: bool, } @@ -1346,14 +1793,17 @@ impl RepartitionStream { schema: Arc, pool: &Arc, inputs: usize, + ordering: SeriesOrdering, ) -> Self { Self { receiver, lease, schema, + ordering, pool: Arc::clone(pool), groups: BTreeMap::new(), eof: Vec::new(), + last_key_from: vec![None; inputs], eof_metadata: None, inputs, output: VecDeque::new(), @@ -1364,10 +1814,13 @@ impl RepartitionStream { } } - fn finish_group(&mut self, key: &ExactIdentity) -> DataFusionResult<()> { + fn finish_group(&mut self, key: &OwnerKey) -> DataFusionResult<()> { let Some(group) = self.groups.remove(key) else { return Ok(()); }; + if matches!(key, OwnerKey::Raw(_)) { + return self.finish_raw_group(group); + } let workspace = if group.fragments.len() > 1 { let rows = group.fragments.iter().fold(0usize, |rows, series| { rows.saturating_add(series.time.len()) @@ -1395,7 +1848,7 @@ impl RepartitionStream { )); }; self.output.push_back(ReadySeries { - series, + series: ReadyOutput::Evaluated(series), charges: ReadyCharges { _charges: group.charges, _metadata: group.metadata, @@ -1406,8 +1859,78 @@ impl RepartitionStream { } } +impl RepartitionStream { + /// Publishes one complete raw identity. + /// + /// Raw pieces are only published once every input reached its normal EOF, because a piece's + /// `is_last` closes that producer's stream, not the series. The workspace this reconstruction + /// needs is charged before it runs and rides with the output until the packer hands the batch + /// off, like any other owner output. + fn finish_raw_group(&mut self, group: PendingGroup) -> DataFusionResult<()> { + let labels = group + .raw_chunks + .first() + .and_then(|chunk| chunk.metric.clone()) + .ok_or_else(|| { + DataFusionError::Execution("raw identity has no explicit labels".to_string()) + })?; + if group + .raw_chunks + .iter() + .any(|chunk| chunk.metric.as_ref() != Some(&labels)) + { + return Err(DataFusionError::Execution( + "raw identity pieces disagree on their explicit labels".to_string(), + )); + } + if group.raw_chunks.len() == 1 { + let mut chunk = group.raw_chunks.into_iter().next().expect("one raw chunk"); + chunk.is_last = true; + self.output.push_back(ReadySeries { + series: ReadyOutput::Raw(chunk), + charges: ReadyCharges { + _charges: group.charges, + _metadata: group.metadata, + _workspace: None, + }, + }); + return Ok(()); + } + // Reconstruction's working set depends on how the pieces interleave, which is only known + // while it runs: charge each allocation as it is made rather than reserving a worst case + // that a well-ordered identity never approaches. + let workspace = MemoryConsumer::new("SeriesRepartitionRawOwner").register(&self.pool); + let mut reserve = |bytes: usize| workspace.try_grow(bytes); + let mut shrink = |bytes: usize| { + let result = workspace.try_shrink(bytes); + debug_assert!( + result.is_ok(), + "raw merge workspace shrink failed: {result:?}" + ); + // The merged data is already correct; tolerate accounting inconsistency rather than + // turning a valid query result into an error. + let _ = result; + }; + let chunk = merge_raw_identity_charging( + &self.schema, + &group.raw_chunks, + &mut reserve, + &mut shrink, + )?; + self.output.push_back(ReadySeries { + series: ReadyOutput::Raw(chunk), + charges: ReadyCharges { + _charges: group.charges, + _metadata: group.metadata, + _workspace: Some(workspace), + }, + }); + Ok(()) + } +} + impl Stream for RepartitionStream { - type Item = Result<(Series, bool), BoxedError>; + type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); @@ -1456,19 +1979,34 @@ impl Stream for RepartitionStream { return Poll::Ready(Some(Err(boxed(DataFusionError::from(&error))))); } if this.rescan { - let next = match &this.resume_after { - Some(after) => this - .groups - .range(( - std::ops::Bound::Excluded(after.clone()), - std::ops::Bound::Unbounded, - )) - .find_map(|(key, group)| { - group_ready(group, &this.eof).then(|| key.clone()) + let next = if this.ordering == SeriesOrdering::KeyThenTime { + this.groups.first_key_value().and_then(|(key, group)| { + group_ready(group, &this.eof, &this.last_key_from, this.ordering, key) + .then(|| key.clone()) + }) + } else { + match &this.resume_after { + Some(after) => this + .groups + .range(( + std::ops::Bound::Excluded(after.clone()), + std::ops::Bound::Unbounded, + )) + .find_map(|(key, group)| { + group_ready( + group, + &this.eof, + &this.last_key_from, + this.ordering, + key, + ) + .then(|| key.clone()) + }), + None => this.groups.iter().find_map(|(key, group)| { + group_ready(group, &this.eof, &this.last_key_from, this.ordering, key) + .then(|| key.clone()) }), - None => this.groups.iter().find_map(|(key, group)| { - group_ready(group, &this.eof).then(|| key.clone()) - }), + } }; match next { Some(key) => { @@ -1493,7 +2031,7 @@ impl Stream for RepartitionStream { .lock() .unwrap_or_else(|e| e.into_inner()) .current = Some(ready.charges); - return Poll::Ready(Some(Ok((ready.series, true)))); + return Poll::Ready(Some(Ok(ready.series))); } if this.reported { return Poll::Ready(None); @@ -1509,36 +2047,67 @@ impl Stream for RepartitionStream { _metadata, })) => { drained += 1; - let Some(first) = chain.first() else { - continue; + let (key, pieces) = match chain { + RoutedChain::Evaluated(chain) => { + let Some(first) = chain.first() else { + continue; + }; + let key = exact_identity(&first.0); + if chain + .iter() + .any(|(series, _, _)| exact_identity(series) != key) + { + this.lease.fail(Arc::new(DataFusionError::Execution( + "one series chain contains multiple exact identities".into(), + ))); + continue; + } + (OwnerKey::Evaluated(key), RoutedChain::Evaluated(chain)) + } + RoutedChain::Raw(chain) => { + let Some(first) = chain.first() else { + continue; + }; + let key = raw_identity(&first.0); + if chain.iter().any(|(chunk, _)| raw_identity(chunk) != key) { + this.lease.fail(Arc::new(DataFusionError::Execution( + "one raw chain contains multiple primary keys".into(), + ))); + continue; + } + if let OwnerKey::Raw(raw_key) = &key + && this.ordering == SeriesOrdering::KeyThenTime + { + let last = &mut this.last_key_from[input]; + if last.as_ref().is_none_or(|last| raw_key > last) { + *last = Some(raw_key.clone()); + } + } + (key, RoutedChain::Raw(chain)) + } }; - let key = exact_identity(&first.0); - if chain - .iter() - .any(|(series, _, _)| exact_identity(series) != key) - { - this.lease.fail(Arc::new(DataFusionError::Execution( - "one series chain contains multiple exact identities".into(), - ))); - continue; - } if !this.groups.contains_key(&key) { let metadata = MemoryConsumer::new("SeriesRepartitionIdentity").register(&this.pool); + let key_bytes = + match &key { + OwnerKey::Evaluated((name, tags)) => { + std::mem::size_of::() + .saturating_add(name.as_ref().map_or(0, |name| name.len())) + .saturating_add(tags.capacity().saturating_mul( + std::mem::size_of::<(Arc, Arc)>(), + )) + .saturating_add( + tags.iter() + .map(|(n, v)| n.len().saturating_add(v.len())) + .sum::(), + ) + } + OwnerKey::Raw(key) => key.capacity(), + }; let bytes = std::mem::size_of::() - .saturating_add(std::mem::size_of::()) - .saturating_add(key.0.as_ref().map_or(0, |name| name.len())) - .saturating_add(key.1.capacity().saturating_mul(std::mem::size_of::<( - Arc, - Arc, - )>( - ))) - .saturating_add( - key.1 - .iter() - .map(|(n, v)| n.len().saturating_add(v.len())) - .sum::(), - ) + .saturating_add(std::mem::size_of::()) + .saturating_add(key_bytes) .saturating_add(this.eof.len().saturating_mul(2)); if let Err(error) = metadata.try_grow(bytes) { this.lease.fail(Arc::new(error)); @@ -1548,6 +2117,7 @@ impl Stream for RepartitionStream { key.clone(), PendingGroup { fragments: Vec::new(), + raw_chunks: Vec::new(), charges: Vec::new(), received: vec![false; this.eof.len()], metadata, @@ -1567,12 +2137,30 @@ impl Stream for RepartitionStream { continue; } group.received[input] = true; - for (series, _, charge) in chain { - group.fragments.push(series); - group.charges.push(charge); + match pieces { + RoutedChain::Evaluated(chain) => { + for (series, _, charge) in chain { + group.fragments.push(series); + group.charges.push(charge); + } + } + RoutedChain::Raw(chain) => { + for (chunk, charge) in chain { + group.raw_chunks.push(chunk); + group.charges.push(charge); + } + } } - if group_ready(group, &this.eof) - && let Err(error) = this.finish_group(&key) + if this.ordering == SeriesOrdering::KeyThenTime { + this.rescan = true; + this.resume_after = None; + } else if group_ready( + group, + &this.eof, + &this.last_key_from, + this.ordering, + &key, + ) && let Err(error) = this.finish_group(&key) { this.lease.fail(Arc::new(error)); } @@ -1612,7 +2200,7 @@ impl Stream for RepartitionStream { .lock() .unwrap_or_else(|e| e.into_inner()) .current = Some(ready.charges); - return Poll::Ready(Some(Ok((ready.series, true)))); + return Poll::Ready(Some(Ok(ready.series))); } return Poll::Ready(None); } @@ -1631,33 +2219,57 @@ fn boxed(error: DataFusionError) -> BoxedError { /// /// The seed is what an execution builds once; a [`XxHash64`] accumulates, so a fresh hasher starts /// from this configuration for each chunk, and that is the only thing each chunk builds. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Debug)] struct Routing { - /// The number of output partitions a hash is reduced modulo. outputs: usize, - /// The seed every chunk of this execution is hashed with. seed: u64, + policy: SeriesRouting, } impl Routing { - /// The routing of an execution with `outputs` output partitions. + /// The canonical routing of an execution with `outputs` output partitions. + #[cfg(test)] fn new(outputs: usize) -> Self { + Self::with_policy(outputs, SeriesRouting::CanonicalLabels) + } + + /// The routing of an execution with `outputs` output partitions and an explicit policy. + fn with_policy(outputs: usize, policy: SeriesRouting) -> Self { Self { outputs, seed: ROUTING_SEED, + policy, } } - /// The output partition `metric`'s canonical label set belongs on. - /// - /// The whole label set is hashed, since that is what a merge groups by: pieces of two series - /// whose label sets differ have to be able to land apart, or the split this node exists to make - /// would not happen. [`label_set`] is what canonicalizes the labels - dropping empty values and - /// keeping the metric name - so two pieces that merge hash the same. + /// The output a label-carrying chunk belongs on. + #[cfg(test)] fn output(&self, metric: &SeriesLabels) -> usize { + self.output_labels(metric) + } + + fn output_labels(&self, metric: &SeriesLabels) -> usize { let mut hasher = XxHash64::with_seed(self.seed); - label_set(metric).hash(&mut hasher); - // `outputs` is at least one, which [`SeriesRepartitionExec::try_new`] is what guarantees. + match &self.policy { + SeriesRouting::CanonicalLabels => label_set(metric).hash(&mut hasher), + SeriesRouting::Matching(policy) => { + crate::vector_match::group_key(metric, policy).hash(&mut hasher) + } + // A raw frame routed by its stored key never reaches here: its destination is the + // stored key itself, not a label hash. Evaluated frames never use this policy. + SeriesRouting::RawPrimaryKey => { + debug_assert!( + false, + "RawPrimaryKey routing hashes stored keys, not labels" + ); + } + } + (hasher.finish() % self.outputs as u64) as usize + } + + fn output_key(&self, key: &[u8]) -> usize { + let mut hasher = XxHash64::with_seed(self.seed); + key.hash(&mut hasher); (hasher.finish() % self.outputs as u64) as usize } } diff --git a/src/promql/src/extension_plan/series_repartition/tests.rs b/src/promql/src/extension_plan/series_repartition/tests.rs index 6dcbff6ef61..46aa8a511cd 100644 --- a/src/promql/src/extension_plan/series_repartition/tests.rs +++ b/src/promql/src/extension_plan/series_repartition/tests.rs @@ -38,7 +38,8 @@ use std::sync::{Arc, Mutex, Weak}; use std::task::{Poll, Waker}; use std::time::Duration; -use datafusion::arrow::array::{Array, Float64Array, ListArray, StructArray}; +use datafusion::arrow::array::{Array, Float64Array, ListArray, StructArray, UInt32Array}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; use datafusion::common::config::ConfigOptions; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -54,6 +55,7 @@ use datatypes::series::Label; use datatypes::series::histogram::{ BucketSideView, CounterResetHint, HistogramBatchBuilder, HistogramView, }; +use datatypes::series::raw::RawTimestamps; use datatypes::series::samples::{FieldSpec, Samples, SeriesFields}; use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; use futures::TryStreamExt; @@ -519,6 +521,236 @@ fn node(child: Arc, outputs: usize) -> Arc Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: "value".into(), + }]), + eval_grid: None, + }) +} + +/// One chunk of primary key `key`, carrying `metric`'s labels, over `times`/`values`. +fn raw_chunk( + key: &str, + metric: Option, + times: &[i64], + values: Vec>, + last: bool, +) -> RawSeriesChunk { + let schema = raw_schema(); + let fields = vec![Samples::from_float(Float64Array::from(values))]; + RawSeriesChunk::try_new( + &schema, + Buffer::from_slice_ref(key.as_bytes()), + metric, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).expect("ascending timestamps"), + SeriesFields::try_from_vec(fields).expect("one sample per timestamp per field"), + last, + ) + .expect("a legal raw chunk") +} + +/// `{__name__="m", host=""}`. +fn raw_metric(host: &str) -> SeriesLabels { + SeriesLabels::try_new( + Some(Arc::from("m")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .expect("legal labels") +} + +/// One raw partition: its chunks, and an optional gate that keeps the stream open afterwards. +#[derive(Clone)] +struct RawPartition { + chunks: Vec, + gate: Option>, + gate_chunks: bool, +} + +struct RawScript { + chunks: std::vec::IntoIter, + schema: Arc, + binding: store_api::region_engine::RawLabelSource, + ordering: SeriesOrdering, + gate: Option>, + gate_chunks: bool, +} + +impl Stream for RawScript { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.gate_chunks + && let Some(gate) = &self.gate + && !gate.is_open() + { + gate.wait(cx); + return Poll::Pending; + } + if let Some(chunk) = self.chunks.next() { + return Poll::Ready(Some(Ok(chunk))); + } + if let Some(gate) = &self.gate + && !gate.is_open() + { + gate.wait(cx); + return Poll::Pending; + } + Poll::Ready(None) + } +} + +impl RawSeriesStream for RawScript { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &self.binding + } + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +/// A child whose partitions hand over raw chunks, one partition per list. +struct RawScripted { + partitions: Vec, + series_schema: Arc, + ordering: SeriesOrdering, + properties: Arc, +} + +impl RawScripted { + fn plan(partitions: Vec) -> Arc { + Self::plan_with_ordering(partitions, SeriesOrdering::Contiguous) + } + + /// The child over `partitions` declaring `ordering` for its keys. + fn plan_with_ordering( + partitions: Vec, + ordering: SeriesOrdering, + ) -> Arc { + let series_schema = raw_schema(); + let frame = common_query::series_frame::raw_series_frame_schema_with_ordering( + &series_schema, + &store_api::region_engine::RawLabelSource::Explicit, + ordering, + ) + .expect("a raw schema declares an explicit-label frame"); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Arc::new(Self { + partitions, + series_schema, + ordering, + properties, + }) + } +} + +impl fmt::Debug for RawScripted { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("RawScripted") + .field("partitions", &self.partitions.len()) + .finish() + } +} + +impl DisplayAs for RawScripted { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "RawScripted: partitions={}", self.partitions.len()) + } +} + +impl ExecutionPlan for RawScripted { + fn name(&self) -> &str { + "RawScripted" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a raw scripted child has no children"); + Ok(self) + } + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let spec = self.partitions.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("raw child has no partition {partition}")) + })?; + pack_raw_series(Box::pin(RawScript { + chunks: spec.chunks.into_iter(), + schema: Arc::clone(&self.series_schema), + binding: store_api::region_engine::RawLabelSource::Explicit, + // The packed batches have to declare what this child's own plan declares: the + // sequence validator and the receiver both read the ordering off the frame, so a + // stream that reported a different one would be validated under the wrong contract. + ordering: self.ordering, + gate: spec.gate, + gate_chunks: spec.gate_chunks, + })) + } +} + +/// The raw owner over `child` routing by an explicit policy. +fn raw_node_with_routing( + child: Arc, + outputs: usize, + routing: SeriesRouting, +) -> DataFusionResult { + SeriesRepartitionExec::try_new_with_routing(child, outputs, routing) +} + +/// The raw owner over `child` grouping by complete stored primary key. +fn raw_node(child: Arc, outputs: usize) -> Arc { + Arc::new( + SeriesRepartitionExec::try_new_with_routing(child, outputs, SeriesRouting::RawPrimaryKey) + .expect("the raw owner is built"), + ) +} + +/// Every chunk of `plan`'s one output, decoded, plus the error it ended with if any. +async fn drain_raw( + plan: &Arc, + context: &Arc, +) -> (Vec, Option) { + let frames = plan + .execute(0, Arc::clone(context)) + .expect("partition 0 exists"); + let mut decoded = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + let mut chunks = Vec::new(); + while let Some(next) = decoded.next().await { + match next { + Ok(chunk) => chunks.push(chunk), + Err(error) => return (chunks, Some(error)), + } + } + (chunks, None) +} + /// Every output partition of `plan`, decoded, and the error it ended with if it ended with one. /// /// All partitions are read at once, as in normal execution. Draining them concurrently makes this @@ -545,6 +777,14 @@ async fn drain( } /// The output partition of `plan` at `partition`, as the frame stream a consumer reads. +/// The evaluated view of one owner publication, for tests that drive the internal stream. +fn evaluated(output: ReadyOutput) -> (Series, bool) { + match output { + ReadyOutput::Evaluated(series) => (series, true), + ReadyOutput::Raw(_) => panic!("an evaluated output carried a raw chunk"), + } +} + fn output( plan: &Arc, partition: usize, @@ -1957,6 +2197,7 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("both identities and EOF are queued", || { producer_stopped(&shared) @@ -1967,8 +2208,9 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { .await .expect("first ready identity") .expect("first identity succeeds"); - assert_eq!(first_metric, first.0.metric); - assert!(first.1); + let (first_series, first_closed) = evaluated(first); + assert_eq!(first_metric, first_series.metric); + assert!(first_closed); assert!( stream.output.is_empty(), "no eagerly merged second result is buffered" @@ -1978,11 +2220,12 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { .await .expect("second ready identity") .expect("second identity succeeds"); - assert_eq!(second_metric, second.0.metric); - assert!(second.1); + let (second_series, second_closed) = evaluated(second); + assert_eq!(second_metric, second_series.metric); + assert!(second_closed); assert!(stream.groups.is_empty()); - drop(first); - drop(second); + drop(first_series); + drop(second_series); assert!( futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) .await @@ -3877,7 +4120,7 @@ async fn a_task_attached_after_the_failure_is_aborted_and_not_counted() { ); let task = tokio::spawn(route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 1, Routing::new(1), Vec::new(), @@ -3959,12 +4202,12 @@ async fn a_recorded_failure_reaches_every_output_once() { .flat_map(|(zero, one)| [zero, one].map(|metric| Step::chunk(chunk(metric, 0, 1, true)))) .collect::>(); steps.push(Step::ErrorSaying(MESSAGE)); + // The producer is gated so it cannot reach its error before this test owns both queues. A + // failure drains every queue nobody has claimed yet, so a producer that ran first would release + // the very charges this test reads, and the retention assertion below could never hold. + let switch = Arc::new(Switch::default()); let plan = node( - Fragile::plan( - vec![Behaviour::OneAtATime(steps)], - Arc::new(Switch::default()), - ) - .plan, + Fragile::plan(vec![Behaviour::Gated(steps)], Arc::clone(&switch)).plan, 2, ); let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); @@ -3976,27 +4219,35 @@ async fn a_recorded_failure_reaches_every_output_once() { let mut streams = Vec::new(); for partition in 0..2 { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); - streams.push(Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )) as SendableEvaluatedSeriesStream); + streams.push(Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }) as SendableEvaluatedSeriesStream); } let execution = shared; + // Both queues are claimed, so the gates below are the only thing between the queued prefix and + // the failure: what is retained now is retained until this test drains it. + assert!( + !switch.is_open(), + "the producer is held before its error, so nothing drained the queues yet" + ); + wait_until("queued output chains retain memory before drain", || { + pool.reserved() > 0 + }) + .await; + switch.open(); wait_until( "the read records its failure while no output is drained", || matches!(execution.ending(), Ending::Failed(_)), ) .await; - // The producer records the failure before its queued chains are all visible, so wait for the - // retention this assertion is about instead of reading the pool once. - wait_until("queued output chains retain memory before drain", || { - pool.reserved() > 0 - }) - .await; let outputs = tokio::time::timeout( Duration::from_secs(10), futures::future::join_all(streams.iter_mut().map(until_end)), @@ -4303,13 +4554,16 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); let frames = pack_evaluated_series( - Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )), + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }), Arc::clone(&series_schema), ) .expect("output frame builds"); @@ -4326,7 +4580,7 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { Poll::Pending }); let peer = route_partition( - Box::pin(peer_stream), + RoutedInput::Evaluated(Box::pin(peer_stream)), 1, Routing::new(2), senders.clone(), @@ -4351,7 +4605,7 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { reported: false, }; let unpolled = route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 0, Routing::new(2), senders.clone(), @@ -4410,13 +4664,16 @@ async fn unpolled_route_drop_preserves_an_already_elected_failure() { for partition in 0..2 { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); let frames = pack_evaluated_series( - Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )), + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }), Arc::clone(&schema()), ) .expect("output frame builds"); @@ -4424,7 +4681,7 @@ async fn unpolled_route_drop_preserves_an_already_elected_failure() { } inputs.failed(1, &Arc::new(DataFusionError::Execution(FIRST.to_string()))); let unpolled = route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 0, Routing::new(2), senders.clone(), @@ -4491,7 +4748,7 @@ async fn ready_chain_yields_after_a_charged_prefix_and_drop_releases_it() { let (sender, receiver) = mpsc::unbounded_channel(); let routing = Routing::new(1); let senders = [sender]; - let mut read = Box::pin(read_chains( + let mut read = Box::pin(read_evaluated_chains( &mut source, 0, &routing, @@ -4567,15 +4824,18 @@ async fn sibling_failure_aborts_budget_yielded_route_and_releases_prefix() { ending: Arc::clone(&ending), }); let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); - let mut output: SendableEvaluatedSeriesStream = Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 1, - )); + let mut output: SendableEvaluatedSeriesStream = Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + SeriesOrdering::Contiguous, + ))), + }); let future = route_partition( - source, + RoutedInput::Evaluated(source), 0, Routing::new(1), vec![sender], @@ -4893,7 +5153,8 @@ async fn discarded_open_chain_eof_is_still_an_error() { Box::pin(futures::stream::iter(vec![Ok((series, false))])); let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); let reservation = MemoryConsumer::new("discarded-open-eof-test").register(&pool); - let result = read_chains(&mut input, 0, &Routing::new(1), &[sender], &reservation).await; + let result = + read_evaluated_chains(&mut input, 0, &Routing::new(1), &[sender], &reservation).await; assert!( matches!(result, Stop::Failed(_)), "open discarded chain EOF is not clean" @@ -5162,6 +5423,7 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("internal producer finishes", || producer_stopped(&shared)).await; // Queue, EOF, and identity-group metadata are charged alongside the retained chunks; the @@ -5181,10 +5443,11 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { .await .unwrap() .unwrap(); - assert!(merged.1, "the logical owner closes once after global EOF"); + let (merged_series, closed) = evaluated(merged); + assert!(closed, "the logical owner closes once after global EOF"); assert_eq!( 4, - merged.0.time.len(), + merged_series.time.len(), "both slices retain their two positions in the merged logical owner" ); assert!( @@ -5197,7 +5460,7 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { .is_none(), "the logical owner has one stable end" ); - drop(merged); + drop(merged_series); drop(stream); assert_eq!( 0, @@ -5457,6 +5720,7 @@ async fn histogram_chain_slices_fit_a_budget_below_the_full_table_charge() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("histogram producer queues the whole chain", || { producer_stopped(&shared) @@ -5483,12 +5747,9 @@ async fn histogram_chain_slices_fit_a_budget_below_the_full_table_charge() { chunks.len(), "one owner series spans all physical slices" ); - assert!(chunks[0].1, "the complete owner closes once"); - assert_eq!( - 256, - chunks[0].0.time.len(), - "all histogram rows survive merge" - ); + let (series, closed) = evaluated(chunks.remove(0)); + assert!(closed, "the complete owner closes once"); + assert_eq!(256, series.time.len(), "all histogram rows survive merge"); assert!(matches!(shared.ending(), Ending::Finished)); assert!( pool.reserved() >= base_bytes + independent_slice_lower_bound, @@ -5763,3 +6024,751 @@ async fn a_chain_that_outgrows_the_pool_ends_every_output_and_releases_it() { }) .await; } + +#[tokio::test] +async fn dropping_a_raw_open_chain_cancels_and_releases_its_reservation() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![ + raw_chunk("k", Some(raw_metric("a")), &[1], vec![Some(1.0)], false), + raw_chunk("k", Some(raw_metric("a")), &[2], vec![Some(2.0)], false), + ], + gate: Some(Arc::clone(&switch)), + // The switch stays closed while this partial chunk is delivered and charged; the producer + // then parks at EOF until the destination is dropped. + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let stream = unpack_raw_series(frames, Some(raw_schema())).unwrap(); + wait_until("raw open chain is reserved", || pool.reserved() > 0).await; + drop(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while !producer_stopped(&execution(&plan)) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the raw stream cancels the open-chain read"); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn dropping_a_raw_destination_releases_its_partial_chain() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + false, + )], + gate: Some(Arc::clone(&switch)), + // The switch stays closed while this partial chunk is delivered and charged; the producer + // then parks at EOF until the destination is dropped. + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let stream = unpack_raw_series(frames, Some(raw_schema())).unwrap(); + wait_until("raw partial chain is reserved", || pool.reserved() > 0).await; + drop(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while !producer_stopped(&execution(&plan)) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the destination stops its raw producer"); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn raw_chain_pool_refusal_releases_partial_reservation() { + let first = raw_chunk("k", Some(raw_metric("a")), &[1], vec![Some(1.0)], false); + let second = raw_chunk("k", Some(raw_metric("a")), &[2], vec![Some(2.0)], false); + // Measure what the first chunk reserves in this exact path, then leave no room for a second. + let probe_pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let probe_switch = Arc::new(Switch::default()); + let probe_plan = raw_node( + RawScripted::plan(vec![RawPartition { + chunks: vec![first.clone()], + gate: Some(Arc::clone(&probe_switch)), + gate_chunks: false, + }]), + 1, + ); + let probe_context = task_context(&probe_pool); + let probe_frames = probe_plan.execute(0, Arc::clone(&probe_context)).unwrap(); + let probe_stream = unpack_raw_series(probe_frames, Some(raw_schema())).unwrap(); + wait_until("probe raw chunk is reserved", || probe_pool.reserved() > 0).await; + let budget = probe_pool.reserved(); + drop(probe_stream); + probe_switch.open(); + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![first, second], + gate: None, + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let context = task_context(&pool); + let (_, error) = tokio::time::timeout(Duration::from_secs(10), drain_raw(&plan, &context)) + .await + .expect("raw pool refusal terminates the execution"); + assert!( + error.is_some(), + "second chunk must be refused after first reservation" + ); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn raw_matching_routing_groups_by_match_key() { + use crate::vector_match::MatchingLabels; + + let with_zone = |zone: &str| { + SeriesLabels::try_new( + Some(Arc::from("m")), + vec![ + Label { + name: Arc::from("host"), + value: Arc::from("h"), + }, + Label { + name: Arc::from("zone"), + value: Arc::from(zone), + }, + ], + ) + .expect("legal labels") + }; + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "ka", + Some(with_zone("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "kb", + Some(with_zone("b")), + &[1], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = Arc::new( + raw_node_with_routing( + child, + 3, + SeriesRouting::Matching(MatchingLabels::On(vec![Arc::from("host")])), + ) + .expect("raw matching routing is accepted"), + ); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let matching = Routing::with_policy( + 3, + SeriesRouting::Matching(MatchingLabels::On(vec![Arc::from("host")])), + ); + let primary = Routing::with_policy(3, SeriesRouting::RawPrimaryKey); + let expected_owner = matching.output_labels(&with_zone("a")); + // Both match keys have host=h, while the primary keys ka/kb route separately. Comparing + // to the positive match-key owner therefore fails if this accidentally routes by primary key. + assert_eq!(expected_owner, matching.output_labels(&with_zone("b"))); + assert_ne!(expected_owner, primary.output_key(b"ka")); + assert_ne!(expected_owner, primary.output_key(b"kb")); + let mut owner = None; + let mut identities = std::collections::HashSet::new(); + for output in 0..plan.outputs() { + let frames = plan + .execute(output, Arc::clone(&context)) + .expect("output exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + while let Some(chunk) = stream.next().await { + let chunk = chunk.expect("matching-routed raw stream succeeds"); + identities.insert(chunk.primary_key.as_slice().to_vec()); + if !chunk.is_last { + continue; + } + if [b"ka".to_vec(), b"kb".to_vec()].contains(&chunk.primary_key.as_slice().to_vec()) { + assert_eq!( + output, expected_owner, + "identity must use its match-key owner" + ); + assert!( + owner + .replace(output) + .is_none_or(|previous| previous == output), + "both match-key identities share an output partition" + ); + } + } + } + assert_eq!( + identities, + [b"ka".to_vec(), b"kb".to_vec()].into_iter().collect() + ); + assert!( + owner.is_some(), + "both identities were published by their common owner" + ); +} + +#[test] +fn raw_canonical_labels_routing_is_refused() { + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![], + gate: None, + gate_chunks: false, + }]); + let error = raw_node_with_routing(child, 1, SeriesRouting::CanonicalLabels) + .expect_err("raw canonical labels routing must remain unsupported"); + assert!(error.to_string().contains("cannot use CanonicalLabels")); +} + +#[tokio::test] +async fn raw_identity_publishes_once_every_input_contributed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1, 3], + vec![Some(1.0), Some(3.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[2, 4], + vec![Some(2.0), Some(4.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan.execute(0, context).expect("partition 0 exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + let mut chunks = Vec::new(); + while let Some(item) = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("published while EOF gate remains closed") + { + chunks.push(item.expect("raw merge succeeds")); + if chunks.last().is_some_and(|chunk| chunk.is_last) { + break; + } + } + let timestamps: Vec = chunks + .iter() + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect(); + assert_eq!(timestamps, [1, 2, 3, 4]); + assert!(chunks.last().is_some_and(|chunk| chunk.is_last)); + switch.open(); +} + +#[tokio::test] +async fn raw_identity_waits_for_an_input_that_has_not_contributed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1, 3], + vec![Some(1.0), Some(3.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[2, 4], + vec![Some(2.0), Some(4.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: true, + }, + ]); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan.execute(0, context).expect("partition 0 exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + assert!( + tokio::time::timeout(Duration::from_millis(200), stream.next()) + .await + .is_err() + ); + switch.open(); + let chunks: Vec<_> = stream.try_collect().await.expect("raw merge succeeds"); + let timestamps: Vec = chunks + .iter() + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect(); + assert_eq!(timestamps, [1, 2, 3, 4]); + assert!(chunks.last().is_some_and(|chunk| chunk.is_last)); +} + +/// Chunks decoded from one frame batch share its Arrow buffers, so charging each chunk the whole +/// allocation would charge the batch's timestamps once per chunk. +/// +/// A second input that stays open and contributes nothing holds every group unpublished, which is +/// what keeps the charges of the first input's chunks alive to be measured. +#[tokio::test] +async fn raw_chunks_of_one_batch_are_charged_their_own_slice() { + const SERIES: usize = 64; + const ROWS: usize = 64; + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: (0..SERIES) + .map(|series| { + let times: Vec = (0..ROWS as i64) + .map(|row| series as i64 * 1000 + row) + .collect(); + let values: Vec> = (0..ROWS).map(|row| Some(row as f64)).collect(); + raw_chunk( + &format!("k{series}"), + Some(raw_metric("a")), + ×, + values, + true, + ) + }) + .collect(), + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(64 << 20)); + let context = task_context(&pool); + let mut frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition 0 exists"); + // Everything the first input has is delivered; the second input has contributed nothing, so no + // group is publishable and every chunk charge of the first input is still held. + // + // Wait for the first charge before sampling: under load the input task may not have run yet, + // and a fixed number of idle ticks would then measure nothing. + wait_until("the first chunk is charged", || pool.reserved() > 0).await; + // Keep the largest charge seen while the rest of the partition is routed, so a charge that + // grows per chunk is measured rather than missed. + let mut peak = 0usize; + let mut idle = 0usize; + while idle < 100 { + let state = futures::future::poll_fn(|cx| Poll::Ready(frames.as_mut().poll_next(cx))).await; + peak = peak.max(pool.reserved()); + match state { + Poll::Ready(Some(Ok(_))) => idle = 0, + Poll::Ready(Some(Err(error))) => panic!("the raw read failed: {error}"), + Poll::Ready(None) => panic!("the read ended before the gated input was opened"), + Poll::Pending => idle += 1, + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + switch.open(); + drop(frames); + // The timestamps of every chunk in the batch come from one buffer of at most SERIES * ROWS + // values. Charging that buffer per chunk would retain about SERIES times its size. + let batch_timestamps = SERIES * ROWS * std::mem::size_of::(); + assert!( + peak < SERIES * batch_timestamps / 2, + "peak {peak} must stay well below the {SERIES}x over-charge of {batch_timestamps} bytes" + ); +} + +/// A histogram field is an index array naming rows of a shared table. A query over few distinct +/// histograms holds mostly indices, and those are charged alongside the table rather than instead +/// of it. +#[tokio::test] +async fn raw_histogram_index_arrays_are_charged() { + const ROWS: usize = 4096; + let side = BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }; + let mut builder = HistogramBatchBuilder::new(); + builder + .push(HistogramView { + schema: 1, + zero_threshold: 0.5, + sum: 1.0, + count: 1.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[], + positive: side, + negative: side, + }) + .expect("one histogram row"); + // Every position names the same row, so the table stays small while the index array grows with + // the number of samples. + let samples = Samples::try_from_histogram( + UInt32Array::from(vec![Some(0u32); ROWS]), + builder.finish().expect("the table is built"), + ) + .expect("one row is referenced by every position"); + let chunk = RawSeriesChunk::try_new( + &raw_schema(), + Buffer::from_slice_ref(b"k"), + Some(raw_metric("a")), + RawTimestamps::try_new(ScalarBuffer::from((0..ROWS as i64).collect::>())) + .expect("ascending timestamps"), + SeriesFields::try_from_vec(vec![samples]).expect("one sample per timestamp"), + true, + ) + .expect("one series' samples are a legal raw chunk"); + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![chunk], + gate: None, + gate_chunks: false, + }, + // An input that has contributed nothing keeps the group unpublished, which is what holds + // the first input's charges alive to be measured. + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(64 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition 0 exists"); + // The chunk is charged as the input task routes it, which happens without this stream being + // polled; wait for the charge instead of sampling it for a fixed number of ticks. + wait_until("the routed chunk is charged", || pool.reserved() > 0).await; + let peak = pool.reserved(); + switch.open(); + drop(frames); + let timestamps = ROWS * std::mem::size_of::(); + let indices = ROWS * std::mem::size_of::(); + assert!( + peak >= timestamps + indices, + "peak {peak} must cover the {timestamps} timestamp bytes and the {indices} index bytes" + ); +} + +/// A key-then-time frame promises its keys arrive in order, so an identity whose inputs have all +/// contributed is published without waiting for any of them to end. +#[tokio::test] +async fn an_ordered_frame_publishes_a_key_its_inputs_have_passed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![ + raw_chunk("a", Some(raw_metric("a")), &[1], vec![Some(1.0)], true), + raw_chunk("b", Some(raw_metric("b")), &[1], vec![Some(2.0)], true), + ], + gate: None, + gate_chunks: false, + }, + // Contributes `a` and then stays open: `a` is complete, and nothing has been seen past + // `b`, so `b` waits. + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[3], + vec![Some(3.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + // A frame may split one identity's rows across chunks, so read until `a` closes. + let mut rows = std::collections::BTreeMap::, Vec>::new(); + let mut closed = std::collections::HashSet::new(); + while !closed.contains(b"a".as_slice()) { + let chunk = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("`a` publishes without EOF") + .expect("`a` arrives") + .expect("valid frame"); + rows.entry(chunk.primary_key.to_vec()) + .or_default() + .extend(chunk.time.as_buffer().iter().copied()); + if chunk.is_last { + closed.insert(chunk.primary_key.to_vec()); + } + } + assert_eq!( + vec![1, 3], + rows[b"a".as_slice()], + "both inputs' samples are merged and `a` closed without any input ending" + ); + assert!( + !rows.contains_key(b"b".as_slice()), + "`b` waits: no input has been seen past it and none has ended" + ); + switch.open(); + while stream.next().await.is_some() {} +} + +/// A watermark lets `a` publish after another input has passed it, without that input ending. +#[tokio::test] +async fn an_ordered_watermark_lets_an_identity_publish_without_input_eof() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + // Delivers `b`, which proves this input passed `a`, then stays open so neither input + // ends before `a` is published. + RawPartition { + chunks: vec![raw_chunk( + "b", + Some(raw_metric("b")), + &[2], + vec![Some(2.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + let first = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("`a` publishes although no input ended") + .expect("`a` arrives") + .expect("valid frame"); + assert_eq!(b"a", first.primary_key.as_slice()); + switch.open(); + while stream.next().await.is_some() {} +} + +/// Without a watermark from the silent input, `a` is not safe to publish until that input ends. +#[tokio::test] +async fn an_ordered_identity_waits_for_an_input_that_has_not_passed_the_key() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: true, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + // The bounded wait itself is the assertion: a publication here would mean the silent input + // was incorrectly treated as having passed `a` (or ended). + assert!( + tokio::time::timeout(Duration::from_millis(100), stream.next()) + .await + .is_err(), + "`a` stays held while another input is silent and open" + ); + switch.open(); + let chunk = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("the opened input lets the output proceed") + .expect("`a` publishes") + .expect("valid frame"); + assert_eq!(b"a", chunk.primary_key.as_slice()); + assert!(chunk.is_last); +} + +#[tokio::test] +async fn raw_identity_groups_by_primary_key_not_labels() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k1", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k2", + Some(raw_metric("a")), + &[2], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (chunks, error) = drain_raw(&plan, &context).await; + assert!( + error.is_none(), + "different keys with identical labels merge independently" + ); + let keys: std::collections::HashSet> = chunks + .iter() + .map(|chunk| chunk.primary_key.as_slice().to_vec()) + .collect(); + assert_eq!(chunks.len(), 2, "one chunk is emitted per primary key"); + assert_eq!(keys, [b"k1".to_vec(), b"k2".to_vec()].into_iter().collect()); +} + +#[tokio::test] +async fn raw_pieces_of_one_key_must_agree_on_labels() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("b")), + &[2], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (_, error) = drain_raw(&plan, &context).await; + let error = error.expect("label disagreement is surfaced"); + assert!( + error + .to_string() + .contains("disagree on their explicit labels") + ); +} + +#[tokio::test] +async fn a_raw_merge_conflict_fails_the_execution() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (_, error) = drain_raw(&plan, &context).await; + let error = error.expect("conflicting duplicate samples fail the execution"); + assert!( + error + .to_string() + .contains("duplicate raw timestamp has conflicting samples") + ); +} diff --git a/src/promql/src/source.rs b/src/promql/src/source.rs index 7d882c7251d..8ad9c086964 100644 --- a/src/promql/src/source.rs +++ b/src/promql/src/source.rs @@ -12,9 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod merge; #[cfg(test)] mod tests; +pub(crate) use merge::merge_raw_identity_charging; + use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; diff --git a/src/promql/src/source/merge.rs b/src/promql/src/source/merge.rs new file mode 100644 index 00000000000..cb0fd0d587a --- /dev/null +++ b/src/promql/src/source/merge.rs @@ -0,0 +1,252 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use datafusion::arrow::buffer::ScalarBuffer; +use datafusion_common::DataFusionError; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::series::histogram::HistogramView; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema}; + +/// Rebuilds one raw series from every piece of one complete physical identity, reporting each +/// allocation before it happens. +/// +/// Reconstruction holds three intermediate vectors whose sizes depend on how the pieces interleave, +/// so what it needs cannot be known up front. The caller charges `reserve` as the work grows +/// instead of reserving a worst case it may never reach. +pub(crate) fn merge_raw_identity_charging( + schema: &SeriesSchema, + chunks: &[RawSeriesChunk], + reserve: &mut dyn FnMut(usize) -> std::result::Result<(), DataFusionError>, + shrink: &mut dyn FnMut(usize), +) -> std::result::Result { + let invalid = |reason: &str| { + DataFusionError::External(Box::new( + InvalidSeriesSnafu { + reason: reason.to_string(), + } + .build(), + )) + }; + let Some(first) = chunks.first() else { + return Err(invalid("cannot merge an empty raw identity")); + }; + if schema.eval_grid.is_some() { + return Err(invalid("raw series schema cannot have an evaluation grid")); + } + let key = first.primary_key.clone(); + let labels = first + .metric + .as_ref() + .ok_or_else(|| invalid("raw identity chunk has no explicit labels"))? + .clone(); + for chunk in chunks { + if chunk.primary_key != key { + return Err(invalid("raw identity chunks have different primary keys")); + } + if chunk.metric.as_ref() != Some(&labels) { + return Err(invalid( + "raw identity chunks have different or missing labels", + )); + } + if chunk.fields.len() != schema.fields.len() + || chunk + .fields + .iter() + .any(|field| field.len() != chunk.time.len()) + { + return Err(invalid("raw identity chunk shape does not match schema")); + } + if chunk + .time + .as_buffer() + .windows(2) + .any(|pair| pair[0] >= pair[1]) + { + return Err(invalid( + "raw identity chunk timestamps are not strictly increasing", + )); + } + } + + let mut row_count = 0usize; + for chunk in chunks.iter() { + row_count = row_count.saturating_add(chunk.time.len()); + } + let chunks_are_ordered_and_disjoint = chunks.windows(2).all(|pair| { + pair[0] + .time + .as_buffer() + .last() + .zip(pair[1].time.as_buffer().first()) + .is_some_and(|(last, first)| last < first) + }); + let mut transient_workspace_bytes = 0usize; + let mut output_fields = Vec::with_capacity(schema.fields.len()); + let times = if chunks_are_ordered_and_disjoint { + reserve(row_count.saturating_mul(std::mem::size_of::()))?; + let mut times = Vec::with_capacity(row_count); + for chunk in chunks { + times.extend_from_slice(chunk.time.as_buffer()); + } + for field_index in 0..schema.fields.len() { + let mut parts = Vec::with_capacity(chunks.len()); + for chunk in chunks { + let source = chunk + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))?; + reserve(std::mem::size_of::())?; + parts.push( + source + .slice(0, chunk.time.len()) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + output_fields.push( + Samples::try_from_chunks(parts) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + times + } else { + transient_workspace_bytes = row_count + .saturating_mul(std::mem::size_of::<(i64, usize, usize)>()) + .saturating_mul(2); + reserve(row_count.saturating_mul(std::mem::size_of::<(i64, usize, usize)>()))?; + let mut rows = Vec::with_capacity(row_count); + for (chunk_index, chunk) in chunks.iter().enumerate() { + for row_index in 0..chunk.time.len() { + rows.push((chunk.time.as_buffer()[row_index], chunk_index, row_index)); + } + } + rows.sort_unstable(); + + reserve( + rows.len() + .saturating_mul(std::mem::size_of::<(i64, usize, usize)>()), + )?; + let mut selected = Vec::with_capacity(rows.len()); + let mut index = 0; + while index < rows.len() { + let mut end = index + 1; + while end < rows.len() && rows[end].0 == rows[index].0 { + end += 1; + } + for duplicate in &rows[index + 1..end] { + for field_index in 0..schema.fields.len() { + if !samples_equal( + chunks[rows[index].1] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))? + .sample(rows[index].2), + chunks[duplicate.1] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))? + .sample(duplicate.2), + ) { + return Err(invalid("duplicate raw timestamp has conflicting samples")); + } + } + } + selected.push(rows[index]); + index = end; + } + + reserve(selected.len().saturating_mul(std::mem::size_of::()))?; + let mut times = Vec::with_capacity(selected.len()); + for row in &selected { + times.push(row.0); + } + for field_index in 0..schema.fields.len() { + let mut parts = Vec::new(); + let mut start = 0; + while start < selected.len() { + let (chunk_index, row_index) = (selected[start].1, selected[start].2); + let mut end = start + 1; + while end < selected.len() + && selected[end].1 == chunk_index + && selected[end].2 == selected[end - 1].2 + 1 + { + end += 1; + } + let source = chunks[chunk_index] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))?; + reserve(std::mem::size_of::())?; + parts.push( + source + .slice(row_index, end - start) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + start = end; + } + output_fields.push( + Samples::try_from_chunks(parts) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + times + }; + let result = RawSeriesChunk::try_new( + schema, + key, + Some(labels), + RawTimestamps::try_new(ScalarBuffer::from(times)) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + SeriesFields::try_from_vec(output_fields) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + true, + ) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + // Only the general merge path allocates transient row-index workspaces. + shrink(transient_workspace_bytes); + Ok(result) +} + +fn samples_equal(left: SampleRef<'_>, right: SampleRef<'_>) -> bool { + match (left, right) { + (SampleRef::Missing, SampleRef::Missing) => true, + (SampleRef::Float(a), SampleRef::Float(b)) => a.to_bits() == b.to_bits(), + (SampleRef::Histogram(a), SampleRef::Histogram(b)) => histograms_equal(a, b), + _ => false, + } +} + +fn histograms_equal(a: HistogramView<'_>, b: HistogramView<'_>) -> bool { + a.schema == b.schema + && a.zero_threshold.to_bits() == b.zero_threshold.to_bits() + && a.sum.to_bits() == b.sum.to_bits() + && a.count.to_bits() == b.count.to_bits() + && a.zero_count.to_bits() == b.zero_count.to_bits() + && a.reset_hint == b.reset_hint + && a.start_timestamp == b.start_timestamp + && float_slices_equal(a.custom_values, b.custom_values) + && a.positive.span_offsets == b.positive.span_offsets + && a.positive.span_lengths == b.positive.span_lengths + && float_slices_equal(a.positive.buckets, b.positive.buckets) + && a.negative.span_offsets == b.negative.span_offsets + && a.negative.span_lengths == b.negative.span_lengths + && float_slices_equal(a.negative.buckets, b.negative.buckets) +} + +fn float_slices_equal(a: &[f64], b: &[f64]) -> bool { + a.len() == b.len() && a.iter().zip(b).all(|(a, b)| a.to_bits() == b.to_bits()) +} diff --git a/src/promql/src/source/merge/tests.rs b/src/promql/src/source/merge/tests.rs new file mode 100644 index 00000000000..1bbc927cbf0 --- /dev/null +++ b/src/promql/src/source/merge/tests.rs @@ -0,0 +1,702 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::arrow::array::{Float64Array, UInt32Array}; +use datafusion::arrow::buffer::Buffer; +use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, SeriesLabels}; + +use super::merge_raw_identity_charging; + +/// The same merge, with accounting the tests do not measure. +fn merge_raw_identity( + schema: &SeriesSchema, + chunks: &[RawSeriesChunk], +) -> datafusion::common::Result { + merge_raw_identity_charging(schema, chunks, &mut |_| Ok(()), &mut |_| {}) +} + +fn schema(count: usize) -> SeriesSchema { + SeriesSchema { + fields: (0..count) + .map(|i| datatypes::series::samples::FieldSpec { + name: Arc::from(format!("f{i}")), + }) + .collect::>() + .into(), + eval_grid: None, + } +} + +fn histogram(sum: f64, threshold: f64) -> Samples { + let mut builder = HistogramBatchBuilder::new(); + let side = BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }; + let index = builder + .push(datatypes::series::histogram::HistogramView { + schema: 1, + zero_threshold: threshold, + sum, + count: 1.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[], + positive: side, + negative: side, + }) + .unwrap(); + Samples::try_from_histogram( + UInt32Array::from(vec![Some(index)]), + builder.finish().unwrap(), + ) + .unwrap() +} + +fn labels(value: &str) -> SeriesLabels { + SeriesLabels::try_new( + Some(Arc::from("metric")), + vec![Label { + name: Arc::from("tag"), + value: Arc::from(value), + }], + ) + .unwrap() +} + +fn chunk( + key: &[u8], + label: &str, + times: Vec, + fields: Vec>>, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(fields.len()), + Buffer::from(key.to_vec()), + Some(labels(label)), + RawTimestamps::try_new(times.into()).unwrap(), + SeriesFields::try_from_vec( + fields + .into_iter() + .map(|values| Samples::from_float(Float64Array::from(values))) + .collect(), + ) + .unwrap(), + false, + ) + .unwrap() +} + +#[test] +fn interleaved_pieces_merge_in_time_order() { + let schema = schema(2); + let result = merge_raw_identity( + &schema, + &[ + chunk( + b"k", + "v", + vec![1, 3], + vec![vec![Some(10.), Some(30.)], vec![Some(1.), Some(3.)]], + ), + chunk( + b"k", + "v", + vec![2, 4], + vec![vec![Some(20.), Some(40.)], vec![Some(2.), Some(4.)]], + ), + ], + ) + .unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3, 4]); + assert_eq!( + result + .fields + .get(0) + .unwrap() + .iter() + .map(|s| match s { + datatypes::series::samples::SampleRef::Float(v) => v, + _ => 0., + }) + .collect::>(), + vec![10., 20., 30., 40.] + ); + assert!(result.is_last); +} + +#[test] +fn ordered_disjoint_pieces_concatenate_without_row_workspace() { + let schema = schema(1); + let chunks = [ + chunk(b"k", "v", vec![1, 2], vec![vec![Some(10.), Some(20.)]]), + chunk(b"k", "v", vec![3, 4], vec![vec![Some(30.), Some(40.)]]), + ]; + let result = merge_raw_identity(&schema, &chunks).unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3, 4]); + assert_eq!(result.time.len(), 4); + assert_eq!( + result + .fields + .get(0) + .unwrap() + .iter() + .map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => 0., + }) + .collect::>(), + vec![10., 20., 30., 40.] + ); + + let mut reserved = 0usize; + let mut shrunk = 0usize; + merge_raw_identity_charging( + &schema, + &chunks, + &mut |bytes| { + reserved += bytes; + Ok(()) + }, + &mut |bytes| shrunk += bytes, + ) + .unwrap(); + assert_eq!( + reserved, + 4 * std::mem::size_of::() + 2 * std::mem::size_of::() + ); + assert_eq!(shrunk, 0); + assert!(shrunk <= reserved); +} + +#[test] +fn touching_ordered_chunks_use_general_duplicate_path() { + let schema = schema(1); + let chunks = [ + chunk(b"k", "v", vec![1, 2], vec![vec![Some(10.), Some(20.)]]), + chunk(b"k", "v", vec![2, 3], vec![vec![Some(20.), Some(30.)]]), + ]; + let mut reserved = 0usize; + let mut shrunk = 0usize; + let result = merge_raw_identity_charging( + &schema, + &chunks, + &mut |bytes| { + reserved += bytes; + Ok(()) + }, + &mut |bytes| shrunk += bytes, + ) + .unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3]); + assert_eq!(shrunk, 2 * 4 * std::mem::size_of::<(i64, usize, usize)>()); + assert!(shrunk <= reserved); +} + +#[test] +fn equal_duplicates_collapse_but_bit_differences_conflict() { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]) + ] + ) + .is_ok() + ); + for (left, right) in [ + (-0.0, 0.0), + ( + f64::from_bits(0x7ff8_0000_0000_0001), + f64::from_bits(0x7ff8_0000_0000_0002), + ), + ] { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![Some(left)]]), + chunk(b"k", "v", vec![1], vec![vec![Some(right)]]) + ] + ) + .is_err() + ); + } + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![None]]), + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]) + ] + ) + .is_err() + ); +} + +#[test] +fn rejects_different_keys_labels_and_internal_duplicates() { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"b", "v", vec![2], vec![vec![Some(2.)]]) + ] + ) + .is_err() + ); + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"a", "other", vec![2], vec![vec![Some(2.)]]) + ] + ) + .is_err() + ); + assert!(RawTimestamps::try_new(vec![1, 1].into()).is_err()); + let mut no_labels = chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]); + no_labels.metric = None; + assert!(merge_raw_identity(&schema(1), &[no_labels]).is_err()); +} + +#[test] +fn histogram_duplicates_compare_every_float_property() { + let schema = schema(1); + let make = |sample| RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(sample), + is_last: false, + }; + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(2.0, 0.5))] + ) + .is_ok() + ); + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(3.0, 0.5))] + ) + .is_err() + ); + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(2.0, -0.0))] + ) + .is_err() + ); +} + +/// Builds a histogram with configurable comparison properties for duplicate-conflict tests. +fn histogram_with(view: datatypes::series::histogram::HistogramView<'_>) -> Samples { + let mut builder = HistogramBatchBuilder::new(); + let index = builder.push(view).unwrap(); + Samples::try_from_histogram( + UInt32Array::from(vec![Some(index)]), + builder.finish().unwrap(), + ) + .unwrap() +} + +/// Two histograms are duplicates only when every property matches. +/// +/// Bucket values are independently free within a valid layout and are compared directly. The +/// span/cursor layout cannot be varied independently of bucket count: each layout must have +/// nonnegative lengths summing to the bucket count and a valid cursor. The cases below vary each +/// offset and length while preserving valid layouts and equal bucket counts. +#[test] +fn histogram_duplicates_compare_every_property() { + let schema = schema(1); + let make = |sample| RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(sample), + is_last: false, + }; + let positive = BucketSideView { + span_offsets: &[0, 0], + span_lengths: &[1, 1], + buckets: &[1.0, 2.0], + }; + let negative = BucketSideView { + span_offsets: &[0, 0], + span_lengths: &[1, 1], + buckets: &[3.0, 4.0], + }; + let build = |schema, + count, + zero_count, + sum, + zero_threshold, + positive, + negative, + reset_hint, + start_timestamp, + custom_values| { + histogram_with(datatypes::series::histogram::HistogramView { + schema, + count, + zero_count, + sum, + zero_threshold, + positive, + negative, + reset_hint, + start_timestamp, + custom_values, + }) + }; + let baseline = || { + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ) + }; + assert!( + merge_raw_identity(&schema, &[make(baseline()), make(baseline())]).is_ok(), + "identical histograms are duplicates" + ); + + let positive_buckets_changed = BucketSideView { + buckets: &[1.0, 2.5], + ..positive + }; + let negative_buckets_changed = BucketSideView { + buckets: &[3.0, 4.5], + ..negative + }; + // Same span count and lengths as the baseline, different cursors: only the offsets differ. + let positive_offsets_changed = BucketSideView { + span_offsets: &[1, 0], + ..positive + }; + let negative_offsets_changed = BucketSideView { + span_offsets: &[1, 0], + ..negative + }; + // Same cursors and same total bucket count (2), different split into spans. A zero-length + // span is a valid layout, so only the lengths differ here. + let positive_lengths_changed = BucketSideView { + span_lengths: &[2, 0], + ..positive + }; + let negative_lengths_changed = BucketSideView { + span_lengths: &[2, 0], + ..negative + }; + let cases: Vec<(&str, Samples)> = vec![ + ( + "schema (1 vs 2)", + build( + 2, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "count (1.0 vs 2.0)", + build( + 1, + 2.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "zero_count (0.0 vs 1.0)", + build( + 1, + 1.0, + 1.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "sum (2.0 vs 3.0)", + build( + 1, + 1.0, + 0.0, + 3.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "zero_threshold (0.5 vs -0.0)", + build( + 1, + 1.0, + 0.0, + 2.0, + -0.0, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive bucket value (2.0 vs 2.5)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_buckets_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative bucket value (4.0 vs 4.5)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_buckets_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive span offset ([0, 0] vs [1, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_offsets_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative span offset ([0, 0] vs [1, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_offsets_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive span lengths ([1, 1] vs [2, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_lengths_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative span lengths ([1, 1] vs [2, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_lengths_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "reset hint (Unknown vs CounterReset)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::CounterReset, + None, + &[], + ), + ), + ( + "start timestamp (None vs Some(7))", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + Some(7), + &[], + ), + ), + ( + "custom values ([] vs [1.5])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[1.5], + ), + ), + ]; + for (what, changed) in cases { + assert!( + merge_raw_identity(&schema, &[make(baseline()), make(changed)]).is_err(), + "two histograms differing only in {what} are a conflict, not a duplicate" + ); + } +} + +#[test] +fn rejects_empty_and_shape_mismatch() { + assert!(merge_raw_identity(&schema(1), &[]).is_err()); + let valid = chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]); + assert!(merge_raw_identity(&schema(2), &[valid]).is_err()); + let malformed = RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.), + Some(2.), + ]))), + is_last: false, + }; + assert!(merge_raw_identity(&schema(1), &[malformed]).is_err()); +} + +#[test] +fn reservation_error_is_preserved() { + let error = merge_raw_identity_charging( + &schema(1), + &[chunk(b"k", "v", vec![1], vec![vec![Some(1.)]])], + &mut |_| { + Err(datafusion::common::DataFusionError::Execution( + "reservation refused".into(), + )) + }, + &mut |_| {}, + ) + .unwrap_err(); + assert!(matches!( + error, + datafusion::common::DataFusionError::Execution(_) + )); +} + +#[test] +fn merge_releases_transient_row_workspace() { + let mut reserved = 0usize; + let mut shrunk = 0usize; + let result = merge_raw_identity_charging( + &schema(1), + &[ + chunk(b"k", "v", vec![1, 3], vec![vec![Some(1.), Some(3.)]]), + chunk(b"k", "v", vec![1, 2], vec![vec![Some(1.), Some(2.)]]), + ], + &mut |bytes| { + reserved = reserved.saturating_add(bytes); + Ok(()) + }, + &mut |bytes| shrunk = shrunk.saturating_add(bytes), + ) + .unwrap(); + assert_eq!(result.time.len(), 3); + assert_eq!(shrunk, 2 * 4 * std::mem::size_of::<(i64, usize, usize)>()); + let retained = reserved - shrunk; + assert!(retained < reserved); + assert!(retained >= result.time.len() * std::mem::size_of::()); +} diff --git a/src/promql/src/vector_match.rs b/src/promql/src/vector_match.rs index f68cc50f454..2196893e3a9 100644 --- a/src/promql/src/vector_match.rs +++ b/src/promql/src/vector_match.rs @@ -1528,9 +1528,9 @@ fn same_grid(lhs: &SeriesSchema, rhs: &SeriesSchema) -> Result> { /// Label pairs identifying a match group or a result; empty values are absent, as in /// Prometheus. -type GroupKey = Vec<(Arc, Arc)>; +pub(crate) type GroupKey = Vec<(Arc, Arc)>; -fn group_key(metric: &SeriesLabels, labels: &MatchingLabels) -> GroupKey { +pub(crate) fn group_key(metric: &SeriesLabels, labels: &MatchingLabels) -> GroupKey { let mut key: GroupKey = match labels { MatchingLabels::On(names) => names .iter() diff --git a/tests-integration/tests/native_prefix.rs b/tests-integration/tests/native_prefix.rs index 9b7685e793c..4cd3b9246db 100644 --- a/tests-integration/tests/native_prefix.rs +++ b/tests-integration/tests/native_prefix.rs @@ -55,13 +55,15 @@ fn evaluated_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef ) } -// Reproduction from the review: 16 hosts × 20k samples, `rate[10s]`, and 16 DN -// partitions produces four batches per chain and fails while unpacking because -// `CoalescePartitionsExec` interleaves their batches, so a chain is cut and the frames no -// longer decode. This test asserts that failure, which is why it runs rather than being -// ignored: it keeps the boundary honest until the frontend region merge becomes -// chain-atomic, and it is the test to change when that fix lands. The parallelism-one -// path is exercised by `native_prefix_executes_over_real_mito_region`. +// Sends `rate[10s]` on a one-point evaluation grid for 16 hosts over 16 region-scan +// partitions to the datanode's remote-read entry point, then checks that the evaluated frames +// decode and each host reports the rate its samples produce. The former decode failure no longer +// reproduces because #10023 makes region package completion close its pending final raw chunk, +// so a chain's value and close are no longer latched apart in separate batches by the +// batch-interleaving `CoalescePartitionsExec`. That coalescing is still not chain-atomic: a chain +// spanning batches can still be cut. This one-point grid gives each host one evaluated value plus +// a possible empty closing chunk, so it does not exercise a chain spanning many batches. The +// parallelism-one path is exercised by `native_prefix_executes_over_real_mito_region`. #[tokio::test(flavor = "multi_thread")] async fn native_prefix_multi_partition_series_chains() { let instance = GreptimeDbStandaloneBuilder::new("native_prefix_partitioned_chains") @@ -205,20 +207,48 @@ async fn native_prefix_multi_partition_series_chains() { } let schema = batches.first().expect("evaluated frames").schema.clone(); let batches = RecordBatches::try_new(schema, batches).unwrap(); - let error = common_query::series_frame::unpack_evaluated_series( + // Package completion closes each region's pending final raw chunk, so the evaluated frames + // decode; ordinary partition coalescing still does not keep chains atomic. Check that every + // host's chain closes exactly once and carries the rate its own samples produce. + let decoded = common_query::series_frame::unpack_evaluated_series( Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())), None, ) .unwrap() .try_collect::>() .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("evaluated series identity changed before its chain closed"), - "{error}" - ); + .expect("a chain-atomic merge decodes every chain"); + + let mut closed = std::collections::BTreeMap::)>::new(); + for (series, is_last) in decoded { + let host = series + .metric + .get("host") + .unwrap_or_else(|| panic!("a decoded series keeps its host label: {:?}", series.metric)) + .to_string(); + let entry = closed.entry(host).or_default(); + if is_last { + entry.0 += 1; + } + for field in series.fields.iter() { + entry + .1 + .extend(field.iter().filter_map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => Some(value), + datatypes::series::samples::SampleRef::Histogram(_) + | datatypes::series::samples::SampleRef::Missing => None, + })); + } + } + assert_eq!(16, closed.len(), "one chain per host: {closed:?}"); + for (host, (closes, values)) in &closed { + assert_eq!(1, *closes, "`{host}` closes its chain exactly once"); + assert_eq!( + vec![1.0], + *values, + "`{host}` reports the rate of its own samples, which a cut chain would lose" + ); + } } #[tokio::test(flavor = "multi_thread")]