diff --git a/docs/formats/mito-samples.md b/docs/formats/mito-samples.md index 4c64823259b..28d6e6fadd7 100644 --- a/docs/formats/mito-samples.md +++ b/docs/formats/mito-samples.md @@ -1,7 +1,22 @@ # Mito columnar samples format This document specifies the version 1 on-disk layout, encodings, ordering rules, -and validation limits for samples artifacts. It shares artifact conventions with +and validation limits for samples artifacts. The version number stays 1: the +exact-identity change is enforced by the mandatory `identity_encoding` manifest +discriminator, not by a version bump. + +Every block names the exact series identity (the canonical sparse primary key), which +the part stores once per series in its [key area](#key-area), followed by that copy's +checksum. +An artifact written before exact identity has the same version number and recorded +only `(table_id, tsid)`, a hash; it can be structurally valid and perfectly ordered +while two distinct series were merged. Nothing in such an artifact proves +otherwise, so a manifest without the `identity_encoding` discriminator is +rejected, and its artifact rebuilt from SSTs; there is no dual reader. The same +applies to an artifact of the earlier layout of this format, whose keys sat inside +the block descriptors: its `identity_encoding` value is the older +`exact-sparse-key`, so it is refused when the manifest is opened rather than when a +part is first read. It shares artifact conventions with [Mito columnar series format](mito-series.md). Samples store nullable Float64 fields, timestamps, sequences, and operation types @@ -51,6 +66,7 @@ generation is unique for each build. A manifest contains: | --- | --- | | `format` | String `columnar-samples` | | `version` | Integer 1 | +| `identity_encoding` | String `exact-sparse-key-area`; required, never defaulted. It names the identity *and* the layout that stores it, so `exact-sparse-key` (keys inline in the block descriptors) is different and not readable | | `region_id` | Physical region ID, u64 | | `source_file_id` | Artifact owner ID matching `SamplesEntry.file_id`; the SST ID for legacy artifacts, or the generation ID for daily artifacts | | `generation` | `FileId` matching the catalog artifact generation | @@ -72,8 +88,14 @@ including the internal names below. Tags are not sample fields and are not store The stored logical Arrow schema is `__table_id: UInt32`, `__tsid: UInt64`, the timestamp column with its recorded unit, `__sequence: UInt64`, `__op_type: UInt8`, -then nullable Float64 fields in manifest order. The first five columns are -non-null. No `__primary_key` or tag columns are stored. +then nullable Float64 fields in manifest order, and finally +`__primary_key: Binary` naming the exact key of the block that produced each row. +The five internal columns and `__primary_key` are non-null; value fields stay +nullable, as the manifest declares them. `__table_id`/`__tsid` are the coarse, +hash-based prefix of that key. Tag columns are not stored; the exact key is the +identity. A reader returns that key once per output batch — each block belongs to +exactly one series — instead of materializing the column per row, and rejects a +request that projects it. Each `files` entry contains `part_id: u32`, `file_size: u64`, `num_rows: u64`, `num_blocks: u64`, `first`, `last`, `min_timestamp: i64`, and `max_timestamp: i64`. @@ -108,10 +130,13 @@ the manifest. │ Block 0: time, sequence, operation, fields │ │ Block 1: ... │ ├──────────────────────────────────────────────┤ ← payload_end +│ Key area: one exact series key per series, │ +│ back to back, not page bounded │ +├──────────────────────────────────────────────┤ ← key_end │ Leaf pages (one is the root at depth 1) │ │ Optional internal root (depth 2) │ ├──────────────────────────────────────────────┤ -│ Trailer 112 bytes │ +│ Trailer 120 bytes │ └──────────────────────────────────────────────┘ ``` @@ -130,8 +155,12 @@ Reserved fields must be zero, and decoded structures must have no trailing bytes ### Column payloads and row semantics -Each block contains exactly one `(table_id, tsid)`. Identity is stored once in -its descriptor. Timestamp, sequence, operation type, and each present value field +Each block contains exactly one series and names it by that series' exact canonical +sparse primary key, which is stored once per series in the part's [key area](#key-area) +and referenced by the block descriptor. The key is the SST's own encoding of +`(table_id, tsid, label column-id/value pairs sorted by name, non-null only)`, and +`(table_id, tsid)` is only its order-preserving prefix. Two distinct series may +share a prefix. Timestamp, sequence, operation type, and each present value field are independent chunks; their references are in the metadata, not inline headers. Emit chunks in that order, with value fields following manifest order. @@ -139,13 +168,17 @@ Rows preserve delete markers, duplicate timestamps, and equal-key ties, includin conflicting values at the same sequence. The format does not encode a deduplicated or snapshot-filtered view. -Identifiers are nondecreasing, timestamps are nondecreasing within an identifier, -and sequences are nonincreasing within equal timestamps, including across parts. -Equal ordering keys retain input order; operation type is not a sorting key. - -`(table_id, tsid)` is the complete identity under the uniqueness assumption. -Equal TSIDs in different tables are distinct. Exact sparse primary keys, tags, -and collision discriminators are not stored. +Block keys are nondecreasing in byte-lexicographic order, timestamps are +nondecreasing within one exact key, and sequences are nonincreasing within equal +timestamps, including across parts. A change of exact key restarts timestamps and +ordinals, even when the coarse `(table_id, tsid)` prefix is unchanged. Equal +ordering keys retain input order; operation type is not a sorting key. The key area +follows the same order, one copy per series. Block keys never decrease, so a key that +is read must not sort before the key read before it; a repeated key must continue its +timestamps and ordinals rather than restart them. Because key bytes are read lazily, +the descriptor-only checks use the `(table_id, tsid)` prefix and the order, disjointness, +and reuse of key references inside one part, and only the keys a request actually loads +are checked for byte order and continuation. Blocks split on an identity change or before the next row would exceed a packing limit. There is no hourly split and no padding of short blocks. The default row @@ -187,6 +220,56 @@ offsets to 4 GiB. Verify the checksum before decompression; require exactly the declared decoded length, with bounded output even if a codec frame claims more. Checksums detect accidental corruption, not malicious modification. +### Key area + +The exact series keys occupy their own region between the column payloads and the +metadata pages: `[payload_end, key_end)`. Keys are written back to back with no +padding and no alignment, in block order, and every block of one series shares one +copy of that series' key. The area therefore holds exactly one key per distinct +series of the part, ordered by exact key. A descriptor's key reference is unique per +series, ordered, and disjoint from every other series' copy, which is why an equal +reference means "the same exact series" without reading the bytes. + +A key's length is bounded by the part, not by a page or a chunk. Version 1 puts no +page inside the key area, so one key may be far longer than the 65,536-byte page +limit; the `u32` length in a descriptor is validated against `key_end`, so a +corrupt length can never address bytes outside the part. + +Each copy is `[key bytes][checksum]`, where the checksum is the u64 XXH64 of those +key bytes with seed zero, stored little-endian. A reference's `length` counts the key +bytes only, so the checksum is always the eight bytes after `offset + length` and is +never part of what the reference names. There is no per-key compression and no other +framing: the key bytes between them are stored exactly as the SST encodes them, and +their own primary-key framing (a fixed reserved prefix plus dense or sparse label +encoding) is what decodes them. The checksum is what protects them. A key previously +sat inside a metadata page, whose stored bytes the page reference already checksummed; +without a per-key checksum, flipping one label byte would keep the framing and the +`(table_id, tsid)` prefix valid and silently attribute the block's rows to another +series. + +Reading stays lazy. A reader reads only the byte ranges of the keys of the blocks a +request selects, at most once per series, and validates a fetched key's checksum, its +framing, and its `(table_id, tsid)` prefix against the descriptor that referenced it. A +reference is rejected as soon as its descriptor is decoded — before any row is emitted +— when it starts before `payload_end`, when `offset + length + 8` exceeds `key_end`, +when the length is zero, or when two consecutive descriptors of one part that do not +name one series overlap. + +A key reference is an offset into one part's key area, so the references of two parts +are unrelated numbers and are never compared. Whether one exact series continues in +another part is decided on the exact key bytes: the last key a request loaded from one +part is compared with the first it loads from the next, and an equal key must continue +timestamps and ordinals, while a greater key starts a new series at ordinal zero. + +Because keys are only read for blocks a request selects, this content check covers the +keys actually loaded: unvisited ranges are validated structurally, not by content. It +is chained, though, so a selection that skips a block, a page, or a whole part cannot +hide the two keys around the gap: the check compares each loaded key with the previous +loaded key, not with the previous traversed block. Loaded keys must be nondecreasing. +Within one exact key the ordinal must strictly increase, and only immediately adjacent +blocks of that same key require it to be exactly one greater; a different key starts a +new series at ordinal zero, so no ordering between the two series' ordinals is required. + ### Metadata pages A leaf describes blocks. An internal page references leaves. The root is either @@ -208,16 +291,20 @@ Leaf (kind 0): Internal root (kind 1): ``` Pages are nonempty and limited to 65,536 decoded bytes, including the 12-byte -prefix. Descriptors never straddle pages. Unlike series label descriptors, -sample descriptors have no variable-size strings, and even the widest supported -descriptor fits in one page. No oversized-leaf exception is needed. +prefix. Descriptors never straddle pages, and every descriptor is fixed-width because +the key it names is not stored in it. The writer requires the 12-byte leaf prefix and +one descriptor to fit in a configured page, which for the 1,024-field maximum is +`12 + 56 + 44 * 1,027 = 45,256` decoded bytes. The page limit is a bound on +descriptors only: it no longer bounds the length of a series key. #### Block descriptor | Field | Bytes | Meaning | | --- | ---: | --- | -| Table ID | 4 | Constant for this block | -| TSID | 8 | Constant for this block | +| Table ID | 4 | Coarse prefix of the constant key in this block | +| TSID | 8 | Coarse prefix of the constant key in this block | +| Key offset | 8 | Absolute offset of the block's exact key in the part's key area | +| Key length | 4 | Nonzero; the key's `(table_id, tsid)` prefix must equal the two fields above | | Minimum timestamp | 8 | Signed i64, first row's timestamp | | Maximum timestamp | 8 | Signed i64, last row's timestamp | | Block ordinal | 8 | Per-identifier ordinal across the artifact | @@ -225,10 +312,14 @@ descriptor fits in one page. No oversized-leaf exception is needed. | Column count | 4 | Number of following column descriptors | | Column descriptors | 44 each | Defined below | -The fixed prefix is 44 bytes. Ordinals start at zero for each identity and -increase by one across block and part boundaries. Blocks sort lexicographically -by `(table_id, tsid, min_timestamp, ordinal)`. Adjacent blocks of the same series -satisfy `previous.max_timestamp <= next.min_timestamp`; equality is permitted. +The fixed prefix is 56 bytes and does not depend on the key: `4 table ID + 8 TSID + +8 key offset + 4 key length + 8 minimum timestamp + 8 maximum timestamp + 8 ordinal + +4 row count + 4 column count`. Ordinals start at zero for each exact key and increase +by one across block and part boundaries while that key continues. Blocks sort by exact +key. Adjacent blocks of the same exact key share one key reference and satisfy +`previous.max_timestamp <= next.min_timestamp`; equality is permitted. Blocks of +different keys have no timestamp relation at all, and a new series' key copy starts at +or after the end of the previous copy. Each column descriptor is: @@ -247,10 +338,15 @@ all rows null is omitted and synthesized as null by the reader; present value chunks have `null_count < row_count`. No value statistics are stored. A block with `P` present fields has `C = 3 + P` chunks and a descriptor of -`44 + 44 * C` bytes. This is 220 bytes for one present field and 45,232 bytes for -1,024 present fields. The leaf prefix makes the latter 45,244 bytes. All chunks -are physically ordered and disjoint in the column area, across blocks as well -as within a block. There is no inline block header or block-wide compression. +`56 + 44 * C` bytes, whatever its key's length. For one present field that is +`56 + 44 * 4 = 232` bytes; at the 1,024-field maximum it is `56 + 44 * 1,027 = 45,244` +bytes, and the 12-byte leaf prefix makes that `45,256`. Key length is therefore not +part of any page budget: a part whose exact key is far longer than a page produces the +same descriptors, and the earlier `base + key_length` ceiling — 65,344 bytes with no +value field, 65,300 with one, and 20,288 at the 1,024-field maximum — no longer +exists. All chunks are physically ordered and disjoint in the column area, across +blocks as well as within a block, and the key area follows the column area. There is no +inline block header or block-wide compression. #### Page reference and time bounds @@ -334,15 +430,17 @@ The series-index format has no sample timestamps or time bounds. | Index depth | 4 | 1 or 2 | | File version | 4 | 1 | | File length | 8 | Includes header and trailer | -| `payload_end` | 8 | End of column area, start of metadata area | +| `payload_end` | 8 | End of column area, start of key area | +| `key_end` | 8 | End of key area, start of metadata area | | Magic | 8 | ASCII `GTSAMPLE` | -| Trailer checksum | 8 | XXH64 of the preceding 104 trailer bytes, seed zero | +| Trailer checksum | 8 | XXH64 of the preceding 112 trailer bytes (the 104-byte body and the magic), seed zero | -The trailer is 112 bytes. Locate it using the recorded object size without a stat +The trailer is 120 bytes. Locate it using the recorded object size without a stat request. Validate the checksum, magic, version, length, depth, area boundaries, root bounds against the manifest part descriptor, and root end position before -traversal. Column references must lie in `[16, payload_end)`; page references lie -in `[payload_end, file_length - 112)`. The header and trailer versions must agree. +traversal. Column references must lie in `[16, payload_end)`; key references must lie +in `[payload_end, key_end)`; page references lie in `[key_end, file_length - 120)`. +The header and trailer versions must agree. ## Encodings and compression @@ -434,7 +532,9 @@ field values. An individual chunk has at most 1 MiB of encoded bytes after decompression and 1 MiB plus 64 KiB of stored bytes. Pages have at most 65,536 decoded and 131,072 -stored bytes. Bound sizes before allocating, fetching, or decompressing. Validate +stored bytes. Series keys are stored in the key area rather than in a page or a +chunk, so a key's length is bounded by the part and by its `u32` reference length, +not by `PAGE_LIMIT`; a key longer than a page is a supported part of this version. Bound sizes before allocating, fetching, or decompressing. Validate counts against enclosing bytes before allocating descriptor vectors; row counts must also satisfy the decoded-block formula. Null output materialization remains subject to the same block bound even when chunks are omitted. @@ -464,12 +564,21 @@ contain several generations with disjoint source sets (at most 64 SSTs per merge and each generation may contain multiple bounded part files. Source Parquet files are never rewritten. -The local series-index cache layout version is 2. On upgrade from layout 1, -startup resets the rebuildable index cache before opening regions; ordinary SST -reads remain available while artifacts rebuild. Rolling back to a layout-1 binary -likewise resets the incompatible cache. Allow for rebuild time and I/O in either -direction; do not reuse daily catalogs with an older reader outside this startup -check. Existing per-SST catalog entries remain readable by the new reader. +This branch has no deployments, so the local series-index cache token is pinned to +`1`: a cache written by any other build is discarded on startup and rebuilt from +the SSTs before regions open, and ordinary SST reads remain available meanwhile. +That is the only cache migration this layout has. An artifact written by an earlier +build of this branch — keys inside the block descriptors, a key area without stored +checksums, a 112-byte trailer without `key_end` — carries a different +`identity_encoding` value, `exact-sparse-key` instead of `exact-sparse-key-area`, so +it is refused when its manifest is opened. Maintenance treats that refusal as "not +reusable" and rebuilds the artifact, so an old artifact is never marked reusable and +then fails on the first part read. A part reached without that check would also be +rejected: the reader locates a trailer from the recorded object size, reads the last +120 bytes of such a part, and refuses the shifted frame instead of decoding it. +Allow for rebuild time and I/O in either direction; do not reuse daily catalogs with +an older reader outside this startup check. Existing per-SST catalog entries remain +readable by the new reader. A prepared scan uses a grouped artifact only when all of its sources belong to that scan's assignment or were excluded solely by its time filter. Partial coverage diff --git a/docs/formats/mito-series.md b/docs/formats/mito-series.md index 5dfe8b46eb5..66ea4bd4017 100644 --- a/docs/formats/mito-series.md +++ b/docs/formats/mito-series.md @@ -1,7 +1,17 @@ # Mito columnar series format This document specifies the version 1 on-disk layout, encodings, ordering rules, -and validation limits for series-index artifacts. For samples, see +and validation limits for series-index artifacts. The version number stays 1: the +exact-identity change is enforced by the mandatory `identity_encoding` manifest +discriminator, not by a version bump. + +Rows are keyed by the exact series identity. A row is one distinct canonical +sparse primary key, so two series that share `(table_id, tsid)` are two rows with +their own labels. An artifact written before exact identity has the same version +number and may contain an already merged row that cannot be told apart from a +complete one; nothing else in the artifact distinguishes them. A manifest that +does not carry `identity_encoding` is therefore rejected, and its artifact rebuilt +from SSTs; there is no dual reader. For samples, see [Mito columnar samples format](mito-samples.md). ## Artifact manifest @@ -25,6 +35,7 @@ One immutable artifact owns the following objects under a unique generation: | --- | --- | | `format` | String `columnar-series` | | `version` | Integer 1 | +| `identity_encoding` | String `exact-sparse-key`; required, never defaulted, and the only thing that separates exact-identity artifacts from older ones | | `region_id` | Physical region ID | | `generation` | Index UUID, using the existing `FileId` serialization | | `tags` | Ordered objects with `column_id` and `name`; all are nullable UTF-8 | @@ -101,11 +112,15 @@ as an independent encoded and compressed chunk, then continue with the next block. Chunks have no inline header: their locations, lengths, codecs, and checksums are recorded in the later [metadata pages](#metadata-pages). -`(__table_id, __tsid)` is the complete identifier **under the current uniqueness -assumption**. Equal TSIDs in different tables are distinct. Records are ordered -lexicographically by this pair and identifiers do not repeat after aggregation. -There is no stored `__primary_key` column. Different label sets may hash to the -same TSID within a table; version 1 stores no collision discriminator. +`(__table_id, __tsid)` is the coarse, hash-based prefix of the identity, not the +identity. Equal TSIDs in different tables are distinct, and different label sets +inside one table may hash to the same TSID. Records are ordered by the canonical +sparse key, which groups equal prefixes together, and a record repeats only for +the same exact key. The block's `first`/`last` identifier bounds stay coarse, so +they overlap for colliding series and only ever over-select. + +There is no stored `__primary_key` column: identity is reconstructed by encoding +the record's own tags in the artifact order, which reproduces the SST key. The row count records input rows aggregated for each identifier. Tags are constant for that identifier under the uniqueness assumption. The count is not necessarily diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index c8b7b68ec77..0a124595252 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1883,6 +1883,127 @@ async fn check_two_phase_series_scan( } } +/// Runs one two-phase series scan over `filters` and returns `(tag_0, field_0, ts)` per row. +async fn two_phase_sparse_rows( + engine: &crate::engine::MitoEngine, + region_id: RegionId, + filters: Vec, +) -> Vec<(String, u64, i64)> { + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + filters, + ..Default::default() + }; + let Scanner::Series(mut scanner) = engine.scanner(region_id, request).await.unwrap() else { + panic!("sparse PerSeries request must use SeriesScan"); + }; + assert_eq!("two_phase", scanner.mode()); + let ranges = scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + scanner + .prepare( + PrepareRequest::default() + .with_ranges(vec![ranges]) + .with_target_partitions(1), + ) + .unwrap(); + let metrics_set = ExecutionPlanMetricsSet::default(); + let context = store_api::region_engine::QueryScanContext::default(); + let batches = scanner + .scan_partition(&context, &metrics_set, 0) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in batches { + let tags = batch.column_by_name("tag_0").unwrap(); + let fields = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + for row in 0..batch.num_rows() { + rows.push(( + datatypes::arrow_array::string_array_value_at_index(tags, row) + .unwrap() + .to_string(), + fields.value(row), + timestamps.value(row), + )); + } + } + rows.sort(); + rows +} + +/// `tsid` is a hash, so two series that share one are still two series. Candidate discovery +/// collapses them into a single identifier, so the data read must apply the tag predicates +/// instead of trusting the assignment as an exact series name. +#[tokio::test] +async fn test_two_phase_series_scan_never_reads_a_colliding_partner() { + let mut env = TestEnv::with_prefix("test_two_phase_colliding_series").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let (metadata, row_schema) = create_two_phase_sparse_region(&engine, region_id).await; + let rows = [ + // One tsid, two label sets. + ( + new_sparse_primary_key(&["host-a", "x"], &metadata, 10, 7), + 1000, + 1, + ), + ( + new_sparse_primary_key(&["host-b", "x"], &metadata, 10, 7), + 2000, + 2, + ), + // A distinct tsid collides with nothing. + ( + new_sparse_primary_key(&["host-c", "x"], &metadata, 10, 8), + 3000, + 3, + ), + ]; + put_sparse_rows(&engine, region_id, &row_schema, rows.to_vec()).await; + test_util::flush_region(&engine, region_id, None).await; + + // Selecting one owner must never return the colliding partner's rows. + assert_eq!( + vec![("host-a".to_string(), 1, 1000)], + two_phase_sparse_rows(&engine, region_id, vec![col("tag_0").eq(lit("host-a"))]).await + ); + // A series without a collision still answers its own selection. + assert_eq!( + vec![("host-c".to_string(), 3, 3000)], + two_phase_sparse_rows(&engine, region_id, vec![col("tag_0").eq(lit("host-c"))]).await + ); + // Without a predicate the path still serves the whole assignment, colliding partner + // included: the fix applies the predicates rather than disabling the scan. + assert_eq!( + vec![ + ("host-a".to_string(), 1, 1000), + ("host-b".to_string(), 2, 2000), + ("host-c".to_string(), 3, 3000), + ], + two_phase_sparse_rows(&engine, region_id, Vec::new()).await + ); + engine.stop().await.unwrap(); +} + /// Scans all partitions in round-robin fashion and returns rows sorted by (tag, ts). /// Also asserts that each series appears in only one partition. async fn collect_partition_rows_round_robin( diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index ef09f5f48ab..6bdc0cf85dc 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -193,9 +193,12 @@ pub struct PrunerOptions { /// /// The flag is baked into the cached [`FileRangeBuilder`], so it belongs to /// the whole scan. All partitions sharing this pruner must agree on it. - /// Two-stage series scans disable this prefilter: candidate discovery applies - /// tag predicates, then the data phase calls [`FileRange::precise_filter_flat`] - /// with tag filtering skipped so the predicates are applied exactly once. + /// Two-stage series scans disable this prefilter because candidate discovery already + /// narrowed the assignment by tag predicate, so running it again would only repeat + /// that work. This is a pruning-hint trade only: the assigned series are + /// `(table_id, tsid)` hashes, so a series that collides with an assigned one shares + /// its identifier without matching the predicate, and the data phase still applies + /// the tag predicates through [`FileRange::precise_filter_flat`]. pub enable_predicate_prefilter: bool, } diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 5ac69d191a6..d95e3cc8005 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -272,9 +272,10 @@ impl SeriesReader { /// Creates a reader for the series assigned to one data partition. /// /// `partition_pruner` must come from the candidate scanner's pruner, which - /// has predicate prefiltering disabled. The file read path applies precise - /// filters with tag filtering skipped because candidate discovery has already - /// enforced tag predicates through the exact assigned-series set. + /// has predicate prefiltering disabled. Candidate discovery narrows the assignment + /// by tag predicate, but an assigned series is a `(table_id, tsid)` hash: a series + /// that collides with it shares the identifier without matching the predicate. The + /// file read path therefore still applies the tag predicates of its precise filter. /// /// A single `range_semaphore` covers both the range-build phase and the final /// merge. @@ -547,10 +548,13 @@ fn scan_series_file_ranges( }; let num_rows_before_filter = record_batch.num_rows(); + // Tag predicates are applied here, not assumed: the assigned series are + // `(table_id, tsid)` hashes, so a colliding partner shares the identifier + // without matching the predicate that selected this series. let Some(record_batch) = range.precise_filter_flat( record_batch, range.pre_filter_mode().skip_fields(), - true, + false, )? else { reader_metrics.filter_metrics.rows_precise_filtered += num_rows_before_filter; diff --git a/src/mito2/src/read/series_scan/prepared.rs b/src/mito2/src/read/series_scan/prepared.rs index 2614cb40ae5..d76303aeabd 100644 --- a/src/mito2/src/read/series_scan/prepared.rs +++ b/src/mito2/src/read/series_scan/prepared.rs @@ -63,8 +63,9 @@ use crate::read::series_scan::{RawScanPreparation, raw}; use crate::read::stream::ScanBatch; use crate::series_index::artifact::ArtifactFormat; use crate::series_index::format::{ - self, FormatReader, SamplesFormatReader, SamplesReadRequest, SeriesReadRequest, + self, FormatReader, SamplesBatch, SamplesFormatReader, SamplesReadRequest, SeriesReadRequest, }; +use crate::series_index::identity::ExactSeriesKey; use crate::series_index::{MetricSeriesId, ROW_COUNT_COLUMN}; fn invalid(reason: impl Into) -> Error { @@ -313,65 +314,16 @@ fn predicate_conjuncts(predicate: Option<&Predicate>) -> Vec<&Expr> { } struct Candidate { + /// Coarse prefix of the exact key; only a hash, never identity. id: MetricSeriesId, + /// The one owned key allocation: the region-qualified exact key. key: Buffer, + /// The canonical sparse key, a view of `key`'s bytes. This is the identity every + /// comparison and sample attribution uses, so it is stored once rather than once per + /// representation. merge_key: Buffer, tags: HashMap>, metric: Option, - metadata_row: Option<(Arc, usize)>, -} - -/// Checks every physical tag before retaining only rows that introduce new identities. -fn compact_candidate_batch( - batch: &RecordBatch, - candidates: &BTreeMap, - tags: &[(u32, String)], -) -> Result { - let table = batch - .column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("invalid table ID column"))?; - let tsid = batch - .column_by_name(DATA_SCHEMA_TSID_COLUMN_NAME) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("invalid TSID column"))?; - let mut new_rows = HashMap::new(); - let mut selected = Vec::new(); - for row in 0..batch.num_rows() { - let id = MetricSeriesId { - table_id: table.value(row), - tsid: tsid.value(row), - }; - let previous = candidates - .get(&id) - .and_then(|c| c.metadata_row.as_ref()) - .map(|(batch, row)| (batch.as_ref(), *row)) - .or_else(|| new_rows.get(&id).map(|row| (batch, *row))); - if let Some((previous, previous_row)) = previous { - for (_, name) in tags { - let current = batch - .column_by_name(name) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("invalid series tag array"))?; - let previous = previous - .column_by_name(name) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("incompatible duplicate label metadata"))?; - if previous.is_null(previous_row) != current.is_null(row) - || (!current.is_null(row) && previous.value(previous_row) != current.value(row)) - { - return Err(invalid("conflicting labels for the same table ID and TSID")); - } - } - } else { - new_rows.insert(id, row); - selected.push(row as u64); - } - } - if selected.len() == batch.num_rows() { - return Ok(batch.clone()); - } - take_record_batch(batch, &UInt64Array::from(selected)).context(ComputeArrowSnafu) } struct ColumnarPreparedSource { @@ -475,7 +427,8 @@ impl ColumnarPreparedSource { preparation_metrics.counts.insert("prepare_batch_series", 0); let direct_labels = sample_schema.is_some(); let conjuncts = predicate_conjuncts(input.predicate_group().predicate()); - let mut candidates = BTreeMap::::new(); + let mut retained_rows = 0_usize; + let mut candidates = BTreeMap::::new(); for id in indexes { let timer = preparation_metrics.timer("columnar_prepare_series_open_wall"); let reader = crate::series_index::columnar::reader::ColumnarReader::open_cached( @@ -554,11 +507,6 @@ impl ColumnarPreparedSource { .or_default() += batch.num_rows() as u64; let started = Instant::now(); let mut duplicate_rows = 0_u64; - if direct_labels { - let rows = batch.num_rows(); - batch = compact_candidate_batch(&batch, &candidates, &tags)?; - duplicate_rows = (rows - batch.num_rows()) as u64; - } let table = batch .column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) .ok_or_else(|| invalid("missing table ID column"))? @@ -583,7 +531,60 @@ impl ColumnarPreparedSource { Ok((*id, strings)) }) .collect::>>()?; - let shared_batch = direct_labels.then(|| Arc::new(batch.clone())); + // Identity is the exact canonical sparse key. `tsid` is a hash, so rows + // that differ only in labels are different series and must both survive: + // re-encoding the artifact's own tags reproduces the SST key byte for byte. + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut offsets = mito_codec::row_converter::SparseOffsetsCache::new(); + let mut identity = Vec::new(); + let mut previous: Option> = None; + let mut fresh = Vec::<(ExactSeriesKey, usize)>::new(); + let mut encoded_key_bytes = 0_u64; + for row in 0..batch.num_rows() { + identity.clear(); + codec + .encode_internal(table.value(row), tsid.value(row), &mut identity) + .map_err(|e| invalid(e.to_string()))?; + codec + .encode_to_vec( + tag_arrays.iter().map(|(id, strings)| { + ( + *id, + if strings.is_null(row) { + ValueRef::Null + } else { + ValueRef::String(strings.value(row)) + }, + ) + }), + &mut identity, + ) + .map_err(|e| invalid(e.to_string()))?; + encoded_key_bytes += identity.len() as u64; + // Adjacent duplicates are the common case and never allocate a key. + if previous.as_deref() == Some(identity.as_slice()) + || candidates.contains_key(identity.as_slice()) + { + duplicate_rows += 1; + continue; + } + let key = ExactSeriesKey::from_sparse_key(&identity, &mut offsets)?; + previous = Some(identity.clone()); + fresh.push((key, row)); + } + // Retain only rows that introduce a new identity; the retained batch backs + // both explicit labels and the metadata rows of direct-label candidates. + let retained = if fresh.len() == batch.num_rows() { + batch.clone() + } else { + let selected: Vec = fresh.iter().map(|(_, row)| *row as u64).collect(); + take_record_batch(&batch, &UInt64Array::from(selected)) + .context(ComputeArrowSnafu)? + }; + if direct_labels { + // Only explicit-label chunks retain a batch to resolve labels from. + retained_rows += retained.num_rows(); + } let label_batch = if direct_labels { let RawLabelSource::PrimaryKey { logical_tags, .. } = &labels else { return Err(invalid("expected bound labels")); @@ -597,7 +598,7 @@ impl ColumnarPreparedSource { .ok_or_else(|| invalid("missing bound tag"))?; Ok(( tag.name.clone(), - batch + retained .schema() .index_of(&name.1) .map_err(|e| invalid(e.to_string()))?, @@ -609,7 +610,7 @@ impl ColumnarPreparedSource { .register(&input.scan_memory_pool), ); let label_batch = - SeriesLabelBatch::try_new(batch.clone(), columns, owner.clone()) + SeriesLabelBatch::try_new(retained.clone(), columns, owner.clone()) .map_err(|e| invalid(e.to_string()))?; owner .try_grow(label_batch.estimated_size()) @@ -618,56 +619,9 @@ impl ColumnarPreparedSource { } else { None }; - let mut encoded_key_bytes = 0_u64; let mut tag_values = 0_u64; let mut batch_series = 0_u64; - for row in 0..batch.num_rows() { - let id = MetricSeriesId { - table_id: table.value(row), - tsid: tsid.value(row), - }; - // The flat merger still needs the compact sparse ID prefix. - // Explicit-label chunks use a region-qualified identity below. - let mut key = if direct_labels { - metadata.region_id.as_u64().to_be_bytes().to_vec() - } else { - Vec::new() - }; - let identity_prefix_len = key.len(); - let codec = SparsePrimaryKeyCodec::schemaless(); - codec - .encode_internal(id.table_id, id.tsid, &mut key) - .map_err(|e| invalid(e.to_string()))?; - let merge_key_len = key.len() - identity_prefix_len; - if !direct_labels { - codec - .encode_to_vec( - tag_arrays.iter().map(|(id, strings)| { - ( - *id, - if strings.is_null(row) { - ValueRef::Null - } else { - ValueRef::String(strings.value(row)) - }, - ) - }), - &mut key, - ) - .map_err(|e| invalid(e.to_string()))?; - encoded_key_bytes += key.len() as u64; - } - if let Some(previous) = candidates.get(&id) { - duplicate_rows += 1; - if previous.key.as_slice() != key.as_slice() { - return Err(invalid( - "conflicting labels for the same table ID and TSID", - )); - } - continue; - } - // Encode borrowed Arrow strings first. Duplicate identities still undergo - // the full label-conflict check, without allocating another owned tag map. + for (retained_row, (exact, row)) in fresh.into_iter().enumerate() { let mut values = HashMap::new(); if !direct_labels { values.reserve(tag_arrays.len()); @@ -679,6 +633,24 @@ impl ColumnarPreparedSource { } } tag_values += values.len() as u64; + // `merge_key` is always the exact key bytes. Explicit-label chunks keep + // the region-qualified form as their `key`, so `primary_key` stays + // region-scoped while identity and ordering stay exact. + let key = if direct_labels { + let mut buffer = Vec::with_capacity(8 + exact.bytes().len()); + buffer.extend_from_slice(&metadata.region_id.as_u64().to_be_bytes()); + buffer.extend_from_slice(exact.bytes()); + Buffer::from(buffer) + } else { + Buffer::from(exact.bytes().to_vec()) + }; + let merge_key = if direct_labels { + key.slice_with_length(8, exact.bytes().len()) + } else { + key.clone() + }; + // The exact key is stored once: `merge_key` is a view of `key`, so the + // reservation charges the single key allocation the candidate retains. let bytes = key.len() + values.values().flatten().map(String::len).sum::() + values.capacity() * (std::mem::size_of::<(u32, Option)>() + 16) @@ -696,24 +668,22 @@ impl ColumnarPreparedSource { SeriesLabels::from_batch( Some(metric_name.clone()), label_batch.clone(), - row, + retained_row, ) .map_err(|e| invalid(e.to_string()))?, ) } else { None }; - let key = Buffer::from(key); - let merge_key = key.slice_with_length(identity_prefix_len, merge_key_len); + let id = exact.id(); candidates.insert( - id, + exact, Candidate { id, key, merge_key, tags: values, metric, - metadata_row: shared_batch.as_ref().map(|batch| (batch.clone(), row)), }, ); } @@ -740,6 +710,9 @@ impl ColumnarPreparedSource { preparation_metrics .counts .insert("prepare_unique_series", candidates.len() as u64); + preparation_metrics + .counts + .insert("prepare_retained_rows", retained_rows as u64); let mut candidates: Vec<_> = candidates.into_values().collect(); candidates.sort_unstable_by(|a, b| a.key.as_slice().cmp(b.key.as_slice())); preparation_metrics.add_elapsed("columnar_prepare_sort", started.elapsed()); @@ -802,48 +775,59 @@ impl PreparedSeriesSource for ColumnarPreparedSource { let sample_schema = sample_schema.clone(); sources.push(Box::pin(async_stream::try_stream! { let reservation = MemoryConsumer::new("ColumnarSamplesRead").register(&context.input.scan_memory_pool); - // Account for the ID selection retained by the samples reader, including - // B-tree node overhead, before allocating it. - let selection_bytes = range.len() * (std::mem::size_of::() + 64); + // Account for the selection retained by the samples reader, including + // B-tree node overhead, before allocating it. The exact-key selection and + // the coarse identifiers the reader derives from it both live for the + // lifetime of the sample stream. + let selection_bytes = candidates[range.clone()] + .iter() + .map(|candidate| { + std::mem::size_of::() + + candidate.merge_key.len() + + std::mem::size_of::() + + 64 + }) + .sum::(); reservation.try_grow(selection_bytes).context(crate::error::ScanMemoryExhaustedSnafu)?; - let ids = candidates[range.clone()].iter().map(|c| c.id).collect::>(); + // Exact ownership: a block is read only when its own key is selected, + // never because a colliding partner shares the identifier. Coarse + // pruning bounds are derived from these keys inside the reader. + // The keys are built from the one key the candidate retains, so the + // selection costs one owned copy per series rather than a second copy + // alongside every candidate. + let mut offsets = mito_codec::row_converter::SparseOffsetsCache::new(); + let keys = candidates[range.clone()] + .iter() + .map(|c| ExactSeriesKey::from_sparse_key(c.merge_key.as_slice(), &mut offsets)) + .collect::>>()?; let mut stage_metrics = crate::metrics::ColumnarReadMetrics::default(); let mut batches = reader.read(SamplesReadRequest { projection: sample_projection(&reader.schema(), &context), memory_pool: Some(context.input.scan_memory_pool.clone()), - ids: Some(ids), + keys: Some(keys), start, end, ..Default::default() })?; - while let Some(batch) = batches.try_next().await? { - let tables = batch.column_by_name(DATA_SCHEMA_TABLE_ID_COLUMN_NAME) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("missing or invalid sample table ID column"))?; - let tsids = batch.column_by_name(DATA_SCHEMA_TSID_COLUMN_NAME) - .and_then(|a| a.as_any().downcast_ref::()) - .ok_or_else(|| invalid("missing or invalid sample TSID column"))?; - let mut offset = 0; - while offset < batch.num_rows() { - let id = MetricSeriesId { table_id: tables.value(offset), tsid: tsids.value(offset) }; - // Sparse keys begin with the order-preserving table ID/TSID prefix. - let ordinal = candidates[range.clone()].binary_search_by_key(&id, |c| c.id) - .map_err(|_| invalid("sample contains an unselected series"))? + range.start; - let mut end = offset + 1; - while end < batch.num_rows() && tables.value(end) == id.table_id && tsids.value(end) == id.tsid { - end += 1; - } - let started = std::time::Instant::now(); - let flat = sample_to_flat(batch.slice(offset, end - offset), &candidates[ordinal], &context, sample_schema.as_ref())?; - let flat = filter_visibility(flat, &context)?; - stage_metrics.add_elapsed("columnar_sample_adapt", started.elapsed()); - offset = end; - if flat.num_rows() == 0 { continue; } - let bytes = flat.get_array_memory_size(); - reservation.try_grow(bytes).context(crate::error::ScanMemoryExhaustedSnafu)?; - yield flat; - reservation.shrink(bytes); - } + while let Some(samples) = batches.try_next().await? { + // A block belongs to exactly one series, so the reader returns its + // exact key once and the whole batch is attributed at once. + let SamplesBatch { batch, key } = samples; + let ordinal = candidates[range.clone()] + .binary_search_by(|candidate| { + candidate.merge_key.as_slice().cmp(key.bytes()) + }) + .map_err(|_| invalid("sample contains an unselected series"))? + + range.start; + let started = std::time::Instant::now(); + let flat = sample_to_flat(batch, &candidates[ordinal], &context, sample_schema.as_ref())?; + let flat = filter_visibility(flat, &context)?; + stage_metrics.add_elapsed("columnar_sample_adapt", started.elapsed()); + if flat.num_rows() == 0 { continue; } + let bytes = flat.get_array_memory_size(); + reservation.try_grow(bytes).context(crate::error::ScanMemoryExhaustedSnafu)?; + yield flat; + reservation.shrink(bytes); } })); } @@ -1122,7 +1106,8 @@ fn filter_visibility(batch: RecordBatch, context: &StreamContext) -> Result, + timestamp: i64, + value: Option, +) { let mut schema = vec![ ColumnSchema { column_name: "__primary_key".into(), @@ -149,7 +168,7 @@ async fn write_with_tags( }, ]; let mut values = vec![ - ValueData::BinaryValue(new_sparse_primary_key(tags, &metadata, 10, tsid)), + ValueData::BinaryValue(primary_key), ValueData::TimestampMillisecondValue(timestamp), ]; let request = if let Some(value) = value { @@ -1063,7 +1082,9 @@ impl SamplesFormatReader for SamplesReadRecorder { fn read( &self, request: SamplesReadRequest, - ) -> Result>> { + ) -> Result< + futures::stream::BoxStream<'static, Result>, + > { self.requests.lock().unwrap().push(request.clone()); self.inner.read(request) } @@ -1162,13 +1183,21 @@ async fn ordinal_range_batches_sample_requests_and_pushes_time_bounds() { { let requests = requests.lock().unwrap(); assert_eq!(source.samples.len(), requests.len()); - let ids = source + // Selection is exact and key-only: the reader derives its coarse bounds from it. + let keys = source .candidates .iter() - .map(|c| c.id) + .map(|c| c.merge_key.to_vec()) .collect::>(); for request in requests.iter() { - assert_eq!(Some(&ids), request.ids.as_ref()); + let requested = request + .keys + .as_ref() + .unwrap() + .iter() + .map(|key| key.bytes().to_vec()) + .collect::>(); + assert_eq!(keys, requested); assert_eq!(Some(1001), request.start); assert_eq!(Some(1601), request.end); } @@ -1206,7 +1235,7 @@ async fn ordinal_range_batches_sample_requests_and_pushes_time_bounds() { .lock() .unwrap() .iter() - .all(|r| r.ids.as_ref().unwrap().len() == 2) + .all(|r| r.keys.as_ref().unwrap().len() == 2) ); engine.stop().await.unwrap(); } @@ -1608,12 +1637,14 @@ async fn metadata_cache_reuses_structures_and_isolates_storage() { let reader = SamplesReader::open_cached(store.clone(), samples.clone(), context) .await .unwrap(); + // The exact key is returned with each batch, so it is never a projected column. let request = SamplesReadRequest { projection: reader .schema() .fields() .iter() .map(|field| field.name().clone()) + .filter(|name| name != store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME) .collect(), ..Default::default() }; @@ -2330,10 +2361,27 @@ async fn overlapping_buckets_retain_only_new_label_rows() { .unwrap(); let prepared = scan.prepared.source.get().unwrap().as_ref().unwrap(); assert_eq!(34, prepared.candidates.len()); - let mut retained = HashMap::new(); + // Duplicate rows from the overlapping bucket are dropped before labels are retained, so no + // series keeps the same timestamp twice. Asserted on this preparation's own rows instead of + // on the process-global counters: concurrent preparations in the same process move those + // counters too, which made a delta comparison fail under parallel test runs. + let chunks = prepared + .read_series(0..prepared.num_series()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + for chunk in &chunks { + let mut seen = std::collections::BTreeSet::new(); + for i in 0..chunk.time.len() { + let timestamp = chunk.time.timestamp(i).unwrap(); + assert!( + seen.insert(timestamp), + "series retained two rows at the same timestamp {timestamp}" + ); + } + } for candidate in prepared.candidates.iter() { - let (batch, _) = candidate.metadata_row.as_ref().unwrap(); - retained.insert(Arc::as_ptr(batch), batch.num_rows()); let expected = if candidate.id.tsid == 0 { "x".to_string() } else { @@ -2344,7 +2392,6 @@ async fn overlapping_buckets_retain_only_new_label_rows() { Some(expected.as_str()) ); } - assert_eq!(34, retained.values().sum::()); engine.stop().await.unwrap(); } @@ -2352,7 +2399,7 @@ async fn overlapping_buckets_retain_only_new_label_rows() { #[case::direct(false)] #[case::mixed_fallback(true)] #[tokio::test] -async fn overlapping_buckets_reject_conflicting_physical_tags(#[case] mixed: bool) { +async fn overlapping_buckets_keep_colliding_physical_tags(#[case] mixed: bool) { let (_env, engine, region) = fixture().await; write_with_tags( &engine, @@ -2385,16 +2432,232 @@ async fn overlapping_buckets_reject_conflicting_physical_tags(#[case] mixed: boo // The conflict is in an unbound physical tag, which must still be checked. *logical_tags = logical_tags[..1].to_vec().into(); } - let error = scan - .prepare_series( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - &binding, - ) + scan.prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding, + ) + .await + .expect("one tsid with two label sets prepares as two series"); + let source = scan.prepared.source.get().unwrap().as_ref().unwrap(); + // Same `(table_id, tsid)`, different physical tags: two series, never one merged row. + assert_eq!(2, source.num_series()); + let keys: HashSet> = (0..source.num_series()) + .map(|ordinal| source.series_info(ordinal).unwrap().primary_key.to_vec()) + .collect(); + assert_eq!(2, keys.len()); + let chunks = source + .read_series(0..source.num_series()) + .unwrap() + .try_collect::>() .await - .err() - .expect("conflicting physical tags must fail preparation"); - assert!(error.output_msg().contains("conflicting labels"), "{error}"); + .unwrap(); + // The differing physical tag is unbound here, so only the exact key separates the + // two candidates; the previous code compared labels and failed preparation. + let identities: HashSet> = source + .candidates + .iter() + .map(|candidate| candidate.merge_key.to_vec()) + .collect(); + assert_eq!(2, identities.len(), "colliding series keep distinct keys"); + if !mixed { + // Every returned chunk names one of the candidates' label sets; a chunk whose + // labels match neither would mean the pair had been merged again. + let labels = |metric: &SeriesLabels| { + metric + .tags() + .iter() + .map(|label| (label.name.to_string(), label.value.to_string())) + .collect::>() + }; + let expected: HashSet> = (0..source.num_series()) + .map(|ordinal| { + labels( + source + .series_info(ordinal) + .unwrap() + .metric + .as_ref() + .unwrap(), + ) + }) + .collect(); + for chunk in &chunks { + let actual = labels(chunk.metric.as_ref().expect("explicit labels")); + assert!( + expected.contains(&actual), + "chunk labels {actual:?} match no candidate: {expected:?}" + ); + } + } + engine.stop().await.unwrap(); +} + +/// Two series forced to share one `tsid` must never receive each other's samples: +/// candidate discovery, sample attribution, and the merged primary key are all exact. +#[tokio::test] +async fn colliding_tsids_isolate_samples_and_labels_per_candidate() { + use common_time::Timestamp; + use common_time::range::TimestampRange; + + let (_env, engine, region) = fixture().await; + let metadata = region.version().metadata.clone(); + let first = new_sparse_primary_key(&["host-020078", "z"], &metadata, 10, 0); + let partner = new_sparse_primary_key(&["host-029838", "z"], &metadata, 10, 0); + // Same `(table_id, tsid)`; only the labels differ. + for (key, base) in [(first, 0.0), (partner, 100.0)] { + for (index, timestamp) in [1000, 2000, 3000].into_iter().enumerate() { + write_primary_key( + &engine, + ®ion, + key.clone(), + timestamp, + Some(base + index as f64), + ) + .await; + } + } + flush_region(&engine, region.region_id, None).await; + build(&engine, ®ion).await; + let make_scan = |predicates: Vec, range: Option| { + let mut input = scan_input( + ®ion, + region.series_index_store.clone().unwrap(), + SelectedFormat::Columnar, + (0..metadata.column_metadatas.len()).collect(), + ); + // Replace the fixture's default predicate, which would hide both hosts. + input.predicate = PredicateGroup::new(&metadata, &predicates).unwrap(); + if let Some(range) = range { + input.time_range = Some(range); + } + SeriesScan::new(input, false) + }; + // One colliding tsid keeps both series, and each ordinal owns only its own rows. + let scan = make_scan(Vec::new(), None); + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let binding = labels(scan.metadata()); + scan.prepare_series(&ctx, &metrics, &binding) + .await + .unwrap() + .unwrap(); + let source = scan.prepared.source.get().unwrap().as_ref().unwrap(); + assert_eq!(3, source.num_series(), "both colliding series must survive"); + let mut value_sets = Vec::new(); + for ordinal in 0..source.num_series() { + let chunks = source + .read_series(ordinal..ordinal + 1) + .unwrap() + .try_collect::>() + .await + .unwrap(); + value_sets.push( + rows(&chunks) + .into_iter() + .map(|(_, _, value)| value) + .collect::>(), + ); + } + assert!(value_sets.contains(&vec![0.0, 1.0, 2.0]), "{value_sets:?}"); + assert!( + value_sets.contains(&vec![100.0, 101.0, 102.0]), + "{value_sets:?}" + ); + assert!( + !value_sets + .iter() + .any(|values| values.contains(&100.0) && values.contains(&0.0)), + "a candidate must not mix both colliding series: {value_sets:?}" + ); + // Filtering one host, inside a time range, returns only that host's samples. + let scan = make_scan( + vec![col("tag_0").eq(lit("host-020078"))], + TimestampRange::new( + Timestamp::new_millisecond(1500), + Timestamp::new_millisecond(2500), + ), + ); + let binding = labels(scan.metadata()); + scan.prepare_series(&ctx, &metrics, &binding) + .await + .unwrap() + .unwrap(); + let source = scan.prepared.source.get().unwrap().as_ref().unwrap(); + assert_eq!(1, source.num_series()); + let chunks = source + .read_series(0..source.num_series()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + rows(&chunks) + .into_iter() + .map(|(_, timestamp, value)| (timestamp, value)) + .collect::>(), + vec![(2000, 1.0)], + "the partner's samples must never be attached to the filtered host" + ); + engine.stop().await.unwrap(); +} + +/// A key past 8 KiB must prepare and read like any other: the series-index writer, the +/// samples writer, and preparation all accept the same key, so no artifact can be published +/// that preparation then refuses. +#[tokio::test] +async fn long_keys_prepare_and_read() { + let (_env, engine, region) = fixture().await; + let metadata = region.version().metadata.clone(); + let label = "x".repeat(8 * 1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &metadata, 10, 0); + assert!(key.len() > 8 * 1024, "{}", key.len()); + for (timestamp, value) in [(1000, 1.0), (2000, 2.0)] { + write_primary_key(&engine, ®ion, key.clone(), timestamp, Some(value)).await; + } + flush_region(&engine, region.region_id, None).await; + build(&engine, ®ion).await; + + let mut input = scan_input( + ®ion, + region.series_index_store.clone().unwrap(), + SelectedFormat::Columnar, + (0..metadata.column_metadatas.len()).collect(), + ); + // The fixture's default predicate would hide this series. + input.predicate = PredicateGroup::new(&metadata, &[]).unwrap(); + let scan = SeriesScan::new(input, false); + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let binding = labels(scan.metadata()); + scan.prepare_series(&ctx, &metrics, &binding) + .await + .unwrap() + .unwrap(); + let source = scan.prepared.source.get().unwrap().as_ref().unwrap(); + // The fixture's own series prepare beside this one; the long key must survive preparation + // byte for byte and own its samples. Served samples make `primary_key` region-qualified. + let mut expected = metadata.region_id.as_u64().to_be_bytes().to_vec(); + expected.extend_from_slice(&key); + let ordinal = (0..source.num_series()) + .find(|ordinal| { + source.series_info(*ordinal).unwrap().primary_key.as_ref() as &[u8] + == expected.as_slice() + }) + .expect("preparation must keep the series whose key is longer than 8 KiB"); + let chunks = source + .read_series(ordinal..ordinal + 1) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + rows(&chunks) + .into_iter() + .map(|(_, timestamp, value)| (timestamp, value)) + .collect::>(), + vec![(1000, 1.0), (2000, 2.0)] + ); engine.stop().await.unwrap(); } diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index c0782893a37..23485394227 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -25,6 +25,7 @@ pub(crate) mod columnar; pub(crate) mod format; #[cfg(test)] mod format_tests; +pub(crate) mod identity; mod maintenance; mod purger; pub(crate) mod samples; diff --git a/src/mito2/src/series_index/columnar.rs b/src/mito2/src/series_index/columnar.rs index 98489924164..904986717f0 100644 --- a/src/mito2/src/series_index/columnar.rs +++ b/src/mito2/src/series_index/columnar.rs @@ -34,6 +34,7 @@ use crate::error::{InvalidMetaSnafu, Result}; use crate::series_index::artifact::ArtifactFormat; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::format::{FormatReader, FormatWriter, SeriesIndexFormat}; +use crate::series_index::identity::IdentityEncoding; use crate::series_index::{ MetricSeriesId, ROW_COUNT_COLUMN, SeriesIndexFileHandle, TABLE_ID_COLUMN, TSID_COLUMN, }; @@ -115,6 +116,8 @@ pub(crate) struct Part { pub(crate) struct Manifest { pub(crate) format: ArtifactFormat, pub(crate) version: u32, + /// Mandatory discriminator; see [`IdentityEncoding`]. Never defaulted on read. + pub(crate) identity_encoding: IdentityEncoding, pub(crate) region_id: u64, pub(crate) generation: FileId, pub(crate) tags: Vec, diff --git a/src/mito2/src/series_index/columnar/metadata.rs b/src/mito2/src/series_index/columnar/metadata.rs index 524c59d44d4..26c79e45360 100644 --- a/src/mito2/src/series_index/columnar/metadata.rs +++ b/src/mito2/src/series_index/columnar/metadata.rs @@ -213,11 +213,11 @@ impl Block { let last = c.u64()?; let rows = c.u32()?; let count = c.u32()? as usize; + // `first`/`last` are the coarse `(table_id, tsid)` prefix, so several rows may + // share one identifier when their exact keys differ. Only their order is fixed, + // so the count is bounded by the block limit and not by the identifier span. require( - rows > 0 - && rows as usize <= codec::COLUMN_LIMIT / 8 - && first <= last - && u64::from(rows - 1) <= last - first, + rows > 0 && rows as usize <= codec::COLUMN_LIMIT / 8 && first <= last, "invalid series block bounds/count", )?; require( @@ -363,8 +363,11 @@ impl Page { for _ in 0..count { let block = Block::decode(&mut c, names, payload_end, &mut end)?; if let Some(previous) = blocks.last() { + // Bounds are the coarse `(table_id, tsid)` prefix and stay + // conservative: two blocks may share one prefix when their exact + // keys differ, so equality is allowed and only ever over-selects. require( - previous.last_id() < block.first_id(), + previous.last_id() <= block.first_id(), "unordered series blocks", )?; } @@ -390,7 +393,10 @@ impl Page { .validate(end, reference.payload.offset, true)?; end = child.payload.end()?; if let Some(previous) = children.last() { - require(previous.last < child.first, "unordered series child ranges")?; + require( + previous.last <= child.first, + "unordered series child ranges", + )?; } children.push(child); } diff --git a/src/mito2/src/series_index/columnar/reader.rs b/src/mito2/src/series_index/columnar/reader.rs index b830dd4339e..0bc4fab233d 100644 --- a/src/mito2/src/series_index/columnar/reader.rs +++ b/src/mito2/src/series_index/columnar/reader.rs @@ -28,7 +28,7 @@ use object_store::ObjectStore; use snafu::ResultExt; use crate::cache::columnar::{self as metadata_cache, MetadataCache}; -use crate::error::{NewRecordBatchSnafu, OpenDalSnafu, RecordBatchSnafu, Result, SerdeJsonSnafu}; +use crate::error::{NewRecordBatchSnafu, OpenDalSnafu, RecordBatchSnafu, Result}; use crate::metrics::ColumnarReadMetrics; use crate::series_index::artifact::ArtifactFormat; use crate::series_index::columnar::codec::{self, Dictionary}; @@ -37,6 +37,7 @@ use crate::series_index::columnar::{ MANIFEST_LIMIT, Manifest, VERSION, invalid, require, unsupported, }; use crate::series_index::format::{FormatReader, SeriesReadRequest}; +use crate::series_index::identity::IdentityEncoding; use crate::series_index::{ MetricSeriesId, ROW_COUNT_COLUMN, SeriesIndexFileHandle, TABLE_ID_COLUMN, TSID_COLUMN, }; @@ -114,10 +115,23 @@ impl ColumnarReader { "series manifest exceeds size limit", )?; let bytes = read_range(&store, &path, 0, length).await?; - let manifest: Manifest = serde_json::from_slice(&bytes).context(SerdeJsonSnafu)?; + // A manifest that the current format cannot read predates exact identity or is + // corrupt; either way it must never be reused. + let manifest: Manifest = serde_json::from_slice(&bytes).map_err(|error| { + invalid(format!( + "series manifest is not readable; rebuild the index: {error}" + )) + })?; if manifest.version != VERSION { return Err(unsupported("series manifest version")); } + // Old artifacts can be structurally valid and perfectly ordered, so identity + // encoding is the only thing that keeps their lost labels out of reads. + if manifest.identity_encoding != IdentityEncoding::CURRENT_SERIES_INDEX { + return Err(unsupported( + "series manifest identity encoding; rebuild the series index", + )); + } require( manifest.format == ArtifactFormat::ColumnarSeries && manifest.region_id == region @@ -155,7 +169,8 @@ impl ColumnarReader { "invalid series part entry", )?; if let Some(last) = previous { - require(last < part.first, "unordered series parts")?; + // Coarse bounds may touch when two colliding series straddle a part. + require(last <= part.first, "unordered series parts")?; } previous = Some(part.last); let part_path = manifest.part_path(part.part_id); @@ -479,10 +494,12 @@ async fn load_column_chunk( let result = if integer { let values = codec::decode_integers(&bytes, rows)?; if column.name == TSID_COLUMN { + // Rows are ordered by exact key, so their coarse TSIDs are non-decreasing: + // a distinct series may repeat its partner's TSID. require( values.first() == Some(&first) && values.last() == Some(&last) - && values.windows(2).all(|v| v[0] < v[1]), + && values.windows(2).all(|v| v[0] <= v[1]), "invalid series TSID ordering/bounds", )?; } else { diff --git a/src/mito2/src/series_index/columnar/tests.rs b/src/mito2/src/series_index/columnar/tests.rs index c9644fdb160..cbaa5402fb4 100644 --- a/src/mito2/src/series_index/columnar/tests.rs +++ b/src/mito2/src/series_index/columnar/tests.rs @@ -298,6 +298,271 @@ async fn aggregation_projection_and_independent_reads() { assert!(reader.read(req).is_err()); } +/// Reads every series row as `(tsid, row count, tag_0)` in artifact order. +async fn read_rows( + format: &ColumnarSeriesFormat, + store: &ObjectStore, + handle: SeriesIndexFileHandle, +) -> Vec<(u64, u64, String)> { + let batches = format + .reader(store.clone(), handle) + .await + .unwrap() + .read(request(&[TSID_COLUMN, ROW_COUNT_COLUMN, "tag_0"])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in &batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let counts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + ids.value(row), + counts.value(row), + labels.value(row).to_owned(), + )); + } + } + rows +} + +/// Publishes a writer that was created from the returned `entry`. +async fn handle_of( + store: &ObjectStore, + mut entry: SeriesIndexEntry, + writer: Box, +) -> SeriesIndexFileHandle { + let artifact = writer.finish().await.unwrap(); + entry.file_size = artifact.file_size; + entry.artifact = Some(artifact); + let (purger, _) = series_index_channel(store.clone()); + SeriesIndexFileHandle::new(metadata().region_id, entry, purger) +} + +/// The captured failure pair: two distinct series with one `table_id`/`tsid` whose +/// timestamps restart. The old writer folded the partner into the first series' row +/// and kept the first series' labels; both rows and both label sets must survive. +#[tokio::test] +async fn colliding_tsids_are_retained_not_merged() { + let store = store(); + let meta = metadata(); + let format = ColumnarSeriesFormat::default(); + let first = new_sparse_primary_key(&["host-020078", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-029838", "z"], &meta, 1, 7); + for dictionary in [false, true] { + let entry = entry(); + let mut writer = format + .writer(meta.clone(), store.clone(), &entry) + .await + .unwrap(); + for timestamp in [100, 200] { + writer + .write(&batch( + &[first.clone(), first.clone()], + dictionary, + timestamp, + )) + .await + .unwrap(); + } + // The partner restarts its timestamps and shares the tsid. + for timestamp in [100, 200] { + writer + .write(&batch( + &[partner.clone(), partner.clone()], + dictionary, + timestamp, + )) + .await + .unwrap(); + } + let handle = handle_of(&store, entry, writer).await; + // Two batches of two rows per series: both rows and both label sets survive, + // and each series keeps its own count. + assert_eq!( + read_rows(&format, &store, handle).await, + vec![ + (7, 4, "host-020078".to_owned()), + (7, 4, "host-029838".to_owned()), + ] + ); + } +} + +/// Many distinct keys under one forced tsid, through both primary-key inputs. +#[tokio::test] +async fn many_distinct_keys_under_one_tsid_are_retained() { + let store = store(); + let meta = metadata(); + let format = ColumnarSeriesFormat::default(); + let keys: Vec<_> = (0..16) + .map(|index| { + let label = format!("host-{index:06}"); + new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 5) + }) + .collect(); + for dictionary in [false, true] { + let entry = entry(); + let mut writer = format + .writer(meta.clone(), store.clone(), &entry) + .await + .unwrap(); + writer.write(&batch(&keys, dictionary, 0)).await.unwrap(); + let handle = handle_of(&store, entry, writer).await; + let rows = read_rows(&format, &store, handle).await; + assert_eq!(rows.len(), keys.len()); + for (index, row) in rows.iter().enumerate() { + assert_eq!(*row, (5, 1, format!("host-{index:06}"))); + } + } +} + +/// Distinct identifiers keep every series, with or without a collision in the input. +#[tokio::test] +async fn distinct_identifiers_keep_every_series() { + let store = store(); + let meta = metadata(); + let format = ColumnarSeriesFormat::default(); + let keys: Vec<_> = (0..16) + .map(|tsid| new_sparse_primary_key(&["a", "x"], &meta, 1, tsid)) + .collect(); + for dictionary in [false, true] { + let handle = build(&format, &store, &[batch(&keys, dictionary, 0)]).await; + let rows = read_rows(&format, &store, handle).await; + assert_eq!(rows.len(), keys.len()); + for (tsid, row) in rows.iter().enumerate() { + assert_eq!(*row, (tsid as u64, 1, "a".to_owned())); + } + } +} + +/// One exact series continuing across input batches, blocks, leaves and parts keeps +/// aggregating into a single row with its labels. +#[tokio::test] +async fn same_exact_series_continues_across_batches_leaves_and_parts() { + let store = store(); + let meta = metadata(); + // One row per block, one descriptor per leaf, two leaves per root. + let format = ColumnarSeriesFormat { + block_rows: 1, + page_bytes: 64, + root_capacity: 2, + compression: Compression::None, + ..Default::default() + }; + let anchor = new_sparse_primary_key(&["a", "x"], &meta, 1, 0); + let rest: Vec<_> = (1..5) + .map(|tsid| new_sparse_primary_key(&["a", "x"], &meta, 1, tsid)) + .collect(); + let mut entry = entry(); + let mut writer = format + .writer(meta.clone(), store.clone(), &entry) + .await + .unwrap(); + for timestamp in [100, 200, 300] { + writer + .write(&batch(&[anchor.clone(), anchor.clone()], false, timestamp)) + .await + .unwrap(); + } + for (index, key) in rest.iter().enumerate() { + writer + .write(&batch(std::slice::from_ref(key), true, 100 + index as i64)) + .await + .unwrap(); + } + let artifact = writer.finish().await.unwrap(); + entry.file_size = artifact.file_size; + entry.artifact = Some(artifact); + let (purger, _) = series_index_channel(store.clone()); + let handle = SeriesIndexFileHandle::new(meta.region_id, entry, purger); + assert!(manifest(&store, &handle).await.files.len() > 1); + let batches = format + .reader(store.clone(), handle) + .await + .unwrap() + .read(request(&[TSID_COLUMN, ROW_COUNT_COLUMN, "tag_0"])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut rows = Vec::new(); + for batch in &batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let counts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + ids.value(row), + counts.value(row), + labels.value(row).to_owned(), + )); + } + } + assert_eq!( + rows[0], + (0, 6, "a".to_owned()), + "the anchor series aggregates across batches" + ); + assert_eq!(rows.len(), 5); + for (index, row) in rows[1..].iter().enumerate() { + assert_eq!(*row, ((index + 1) as u64, 1, "a".to_owned())); + } +} + +/// Regressions in sparse-key order are refused instead of opening a second row for +/// a series that was already emitted. +#[tokio::test] +async fn out_of_order_series_are_refused() { + let store = store(); + let meta = metadata(); + let format = ColumnarSeriesFormat::default(); + let first = new_sparse_primary_key(&["a", "x"], &meta, 1, 1); + let second = new_sparse_primary_key(&["a", "x"], &meta, 1, 2); + let mut writer = format + .writer(meta.clone(), store.clone(), &entry()) + .await + .unwrap(); + writer + .write(&batch(&[first.clone(), second], false, 0)) + .await + .unwrap(); + let error = writer + .write(&batch(&[first], false, 1)) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("not sorted by primary key"), "{error}"); + writer.abort().await.unwrap(); +} + #[tokio::test] async fn inconsistent_object_sizes_are_rejected() { let store = store(); diff --git a/src/mito2/src/series_index/columnar/writer.rs b/src/mito2/src/series_index/columnar/writer.rs index 05fe94ee81b..2a9c36238d6 100644 --- a/src/mito2/src/series_index/columnar/writer.rs +++ b/src/mito2/src/series_index/columnar/writer.rs @@ -35,6 +35,7 @@ use crate::series_index::columnar::{ ColumnarSeriesFormat, MANIFEST_LIMIT, Manifest, Part, Tag, VERSION, invalid, require, }; use crate::series_index::format::FormatWriter; +use crate::series_index::identity::{ExactSeriesKey, IdentityEncoding}; use crate::series_index::writer::{decode_series_identity, tag_columns, validate_metadata}; use crate::series_index::{MetricSeriesId, ROW_COUNT_COLUMN, TSID_COLUMN}; @@ -43,7 +44,8 @@ const BLOCK_BYTE_BUDGET: usize = 8 * 1024 * 1024; #[derive(Debug)] struct Row { - id: MetricSeriesId, + /// The exact series identity; `id` is only its prefix. + key: ExactSeriesKey, count: u64, tags: Vec>, } @@ -89,6 +91,7 @@ impl ColumnarWriter { let manifest = Manifest { format: ArtifactFormat::ColumnarSeries, version: VERSION, + identity_encoding: IdentityEncoding::CURRENT_SERIES_INDEX, region_id: metadata.region_id.as_u64(), generation: entry.index_uuid, tags: tags @@ -140,7 +143,7 @@ impl ColumnarWriter { || self .rows .first() - .is_some_and(|first| first.id.table_id != row.id.table_id) + .is_some_and(|first| first.key.id().table_id != row.key.id().table_id) { self.flush_block().await?; } @@ -153,8 +156,14 @@ impl ColumnarWriter { Ok(()) } + /// Aggregates one row of input. + /// + /// Identity is the exact sparse primary key, never `(table_id, tsid)`: distinct + /// series may share a `tsid`, and both keep their own row and labels. Rows carry + /// their tags and are ordered by exact key, so the format can represent the pair + /// without inventing a second identity. async fn append_key(&mut self, key: &[u8]) -> Result<()> { - if key == self.last_key && self.current.is_some() { + if self.current.is_some() && key == self.last_key.as_slice() { if let Some(row) = self.current.as_mut() { row.count = row .count @@ -163,24 +172,26 @@ impl ColumnarWriter { } return Ok(()); } - let (id, tags) = + let (exact, tags) = decode_series_identity(key, &self.tags, &mut self.offsets, &mut self.buffer)?; - if let Some(row) = self.current.as_mut() { - require(id >= row.id, "series identifiers are not sorted")?; - if id == row.id { - row.count = row - .count - .checked_add(1) - .ok_or_else(|| invalid("series row count overflow"))?; - self.last_key.clear(); - self.last_key.extend(key); - return Ok(()); - } + if self.current.is_some() { + // Input arrives in sparse-key order. A regression is a contract violation, + // and the previous key cannot be the current one: equal keys return above. + // A different key that shares `(table_id, tsid)` is a different series and + // becomes its own row. + require( + key > self.last_key.as_slice(), + "series input is not sorted by primary key", + )?; } self.push_current().await?; - self.current = Some(Row { id, count: 1, tags }); + self.current = Some(Row { + key: exact, + count: 1, + tags, + }); self.last_key.clear(); - self.last_key.extend(key); + self.last_key.extend_from_slice(key); Ok(()) } @@ -421,8 +432,13 @@ fn encode_block( let first = rows .first() .ok_or_else(|| invalid("empty series block"))? - .id; - let last = rows.last().ok_or_else(|| invalid("empty series block"))?.id; + .key + .id(); + let last = rows + .last() + .ok_or_else(|| invalid("empty series block"))? + .key + .id(); let mut block = Block { table_id: first.table_id, first: first.tsid, @@ -434,7 +450,7 @@ fn encode_block( for (name, values) in [ ( TSID_COLUMN, - rows.iter().map(|r| r.id.tsid).collect::>(), + rows.iter().map(|r| r.key.id().tsid).collect::>(), ), (ROW_COUNT_COLUMN, rows.iter().map(|r| r.count).collect()), ] { diff --git a/src/mito2/src/series_index/format.rs b/src/mito2/src/series_index/format.rs index f2904bd8897..8dd3aad9aac 100644 --- a/src/mito2/src/series_index/format.rs +++ b/src/mito2/src/series_index/format.rs @@ -33,6 +33,7 @@ use store_api::storage::consts::ReservedColumnId; use crate::error::{InvalidRecordBatchSnafu, Result}; use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata, SamplesFileHandle}; use crate::series_index::catalog::SeriesIndexEntry; +use crate::series_index::identity::ExactSeriesKey; use crate::series_index::{MetricSeriesId, SeriesIndexFileHandle}; use crate::sst::PARQUET_FIELD_ID_KEY; @@ -89,14 +90,37 @@ pub(crate) struct SamplesReadRequest { pub(crate) memory_pool: Option>, pub(crate) projection: Vec, pub(crate) table_id: Option, - pub(crate) ids: Option>, + /// Exact selection. A block whose key is absent here is never read, even when its + /// coarse `(table_id, tsid)` is present. The reader derives its coarse pruning set + /// from these keys, so the selection cannot disagree with itself. + pub(crate) keys: Option>, pub(crate) start: Option, pub(crate) end: Option, } +/// One samples output batch and the exact series that owns every row in it. +/// +/// A sample block belongs to exactly one series, so the exact key travels with the batch +/// instead of being repeated once per row: consumers attribute the whole batch at once. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct SamplesBatch { + pub(crate) batch: RecordBatch, + pub(crate) key: ExactSeriesKey, +} + +/// Batches read as batches: the key is the only column the wrapped batch does not carry. +impl std::ops::Deref for SamplesBatch { + type Target = RecordBatch; + + fn deref(&self) -> &RecordBatch { + &self.batch + } +} + pub(crate) trait SamplesFormatReader: Send + Sync { fn schema(&self) -> SchemaRef; - fn read(&self, request: SamplesReadRequest) -> Result>>; + fn read(&self, request: SamplesReadRequest) + -> Result>>; } /// Builds and opens immutable per-SST samples artifacts. @@ -159,13 +183,23 @@ pub(crate) fn field_id(field: &Field) -> Result { /// Checks whether samples can be read with the current region schema. pub(crate) fn validate_sample_schema(schema: &SchemaRef, metadata: &RegionMetadata) -> Result<()> { + // Exact identity: readers attribute every sample row by this column, so an + // artifact that cannot expose it is unusable. + let key = schema + .field_with_name("__primary_key") + .map_err(|_| invalid("samples lack the exact series key column; rebuild series indexes"))?; + if *key.data_type() != datatypes::arrow::datatypes::DataType::Binary { + return Err(invalid( + "samples series key column is not binary; rebuild series indexes", + )); + } let mut ids = HashSet::new(); - for field in schema - .fields() - .iter() - .skip(2) - .filter(|f| f.name() != "__sequence" && f.name() != "__op_type") - { + for field in schema.fields().iter().skip(2).filter(|f| { + !matches!( + f.name().as_str(), + "__sequence" | "__op_type" | "__primary_key" + ) + }) { let id = field_id(field)?; let column = metadata .column_by_id(id) @@ -260,6 +294,10 @@ mod tests { ) }), ); + if !tags { + // Current artifacts expose the exact series key of each row. + fields.push(Field::new("__primary_key", DataType::Binary, false)); + } let validate = |fields: Vec| { let schema = Arc::new(Schema::new(fields)); if tags { diff --git a/src/mito2/src/series_index/identity.rs b/src/mito2/src/series_index/identity.rs new file mode 100644 index 00000000000..a1acc2c16fa --- /dev/null +++ b/src/mito2/src/series_index/identity.rs @@ -0,0 +1,384 @@ +// 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. + +//! The exact series identity that SSTs already store. +//! +//! `tsid` is a hash, so two distinct series can share one [`MetricSeriesId`]. The +//! sparse primary key stored by the SST is exact, so every place that decides +//! whether two rows belong to the same series decides it on [`ExactSeriesKey`]. + +use std::cmp::Ordering; +use std::fmt; +use std::hash::{Hash, Hasher}; + +use mito_codec::row_converter::SparseOffsetsCache; +use mito_codec::row_converter::sparse::SparsePrimaryKeyView; +use snafu::ResultExt; +use store_api::storage::ColumnId; + +use crate::error::{DecodeSnafu, Result}; +use crate::series_index::MetricSeriesId; + +/// The identity semantics *and storage layout* an artifact was written with. +/// +/// This is the one mandatory discriminator every artifact manifest carries. Artifacts +/// written under hash-only identity are structurally valid and perfectly ordered, so +/// nothing else in them distinguishes them from exact ones: an already collapsed series +/// index looks fine and its lost labels cannot be recovered. A manifest that does not +/// name this encoding is therefore not reusable, whatever its coverage says. +/// +/// The serialized value names the layout as well as the identity, because two artifacts can +/// store the same exact key in incompatible places: the samples format kept one copy per block +/// inside its pages, then moved to one copy per series in a dedicated key area with a stored +/// checksum. An artifact of the older layout is structurally valid and ordered too, so without +/// the layout in this value maintenance would mark it reusable and only a part read would fail. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub(crate) enum IdentityEncoding { + /// Identity is [`ExactSeriesKey`], stored inside the artifact's own pages: the + /// series-index layout, and the retired samples layout that repeated the key in every + /// block descriptor. + #[serde(rename = "exact-sparse-key")] + ExactSparseKey, + /// Identity is [`ExactSeriesKey`], stored once per series in a part-level key area after + /// the column payloads, each copy followed by its checksum. + #[serde(rename = "exact-sparse-key-area")] + ExactSparseKeyArea, +} + +impl IdentityEncoding { + /// The encoding new samples artifacts are written with; their readers require exactly this. + /// + /// Distinct from the series-index encoding on purpose: it names the key-area layout, so an + /// artifact written before the key area cannot be taken for a current one. + pub(crate) const CURRENT_SAMPLES: Self = Self::ExactSparseKeyArea; + + /// The encoding new series-index artifacts are written with; their readers require this. + pub(crate) const CURRENT_SERIES_INDEX: Self = Self::ExactSparseKey; +} + +/// The exact identity of one series: its canonical sparse primary-key bytes. +/// +/// The key is the SST's own encoding of +/// `(table_id, tsid, label column-id/value pairs sorted by name, non-null only)`, +/// as produced by [`mito_codec::row_converter::sparse::SparsePrimaryKeyCodec`]. No +/// second label encoding exists: callers hand in the bytes they read from the SST. +/// +/// # Ordering +/// +/// Identity and order are both defined on the encoded bytes only: +/// +/// * `PartialEq`/`Hash` use the bytes, so two keys are the same series iff their +/// encodings are equal. +/// * `Ord` is byte-lexicographic. The reserved `table_id`/`tsid` prefix is fixed +/// width and order preserving, so this order groups every key that shares one +/// [`MetricSeriesId`] together and sorts those groups by `MetricSeriesId`. +/// Inside one group it orders by the encoded `(column_id, value)` label sequence. +/// +/// This is also the order SSTs store rows in, so artifact writers receive input in +/// key order and can reject regressions instead of guessing. +/// +/// `id` is always derived from `key` and never participates in comparison: every +/// comparison below reads `key` only. +/// +/// # Assumptions +/// +/// This is *storage* identity, not Prometheus label-set identity. It relies on the +/// existing sparse format's invariants, none of which it redefines: +/// +/// * `tsid` is a stable hash for a label set within one logical table; a rebuild of +/// the same rows writes the same `tsid`. +/// * Column IDs are stable for a region, so the encoded `(column_id, value)` pairs +/// identify the same labels over time. Re-creating a region with different IDs +/// changes the key bytes. +/// * Labels are encoded in the codec's canonical order (sorted by column name), and +/// null labels are omitted while empty strings are encoded. This type never +/// normalizes either case. +/// +/// Equal keys therefore mean "the same series as far as this database stores it". +/// They are not claimed to be equal to any Prometheus series identity. +#[derive(Clone)] +pub(crate) struct ExactSeriesKey { + /// Order-preserving `(table_id, tsid)` prefix decoded from `key`. + id: MetricSeriesId, + /// Canonical sparse primary-key encoding of the whole series. + key: Box<[u8]>, +} + +impl ExactSeriesKey { + /// Wraps one encoded sparse primary key after validating its framing. + /// + /// # Length + /// + /// This entry point deliberately has no length cap, and neither does any other: + /// key length is bounded by the containers that store a key (one key per block, + /// plus the descriptor, page, and column limits), and a cap enforced on only one + /// entry point would let a writer publish a key another one refuses. Long labels + /// are real data — a single 8 KiB label already encodes to a larger sparse key — + /// so framing validation is the only requirement here. + pub(crate) fn from_sparse_key(key: &[u8], offsets: &mut SparseOffsetsCache) -> Result { + let mut view = SparsePrimaryKeyView::new(key, offsets).context(DecodeSnafu)?; + // `new` validates the reserved prefix only; walk the labels so a truncated or + // trailing-garbage key cannot become an identity that happens to share a prefix. + view.visit_label_ids(|_| Ok(())).context(DecodeSnafu)?; + Ok(Self::from_view(&view, key)) + } + + /// The reserved `(table_id, tsid)` prefix. It is a coarse bucket, not identity. + pub(crate) fn id(&self) -> MetricSeriesId { + self.id + } + + /// The canonical bytes; the only thing that decides identity. + pub(crate) fn bytes(&self) -> &[u8] { + &self.key + } + + /// Builds the key from a view already validated against `key`. + fn from_view(view: &SparsePrimaryKeyView<'_, '_>, key: &[u8]) -> Self { + Self { + id: MetricSeriesId { + table_id: view.table_id(), + tsid: view.tsid(), + }, + key: key.into(), + } + } +} + +// Only `key` participates in identity. `id` is a prefix function of `key`, so a +// future field cannot accidentally join (or reorder) the identity. + +impl PartialEq for ExactSeriesKey { + fn eq(&self, other: &Self) -> bool { + self.key == other.key + } +} + +impl Eq for ExactSeriesKey {} + +impl PartialOrd for ExactSeriesKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ExactSeriesKey { + fn cmp(&self, other: &Self) -> Ordering { + self.key.cmp(&other.key) + } +} + +impl fmt::Debug for ExactSeriesKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ExactSeriesKey") + .field("id", &self.id) + .field("key", &Hex(&self.key)) + .finish() + } +} + +/// Lets a map keyed by the exact key be probed with the raw canonical bytes it was +/// built from, so hot loops do not allocate a key per row. `Ord`/`Eq` are byte +/// comparisons, which is exactly `[u8]`'s order, so the borrow is consistent. +impl std::borrow::Borrow<[u8]> for ExactSeriesKey { + fn borrow(&self) -> &[u8] { + &self.key + } +} + +impl Hash for ExactSeriesKey { + fn hash(&self, state: &mut H) { + self.key.hash(state); + } +} + +/// Prints at most [`HEX_LIMIT`] bytes of a key, so logs stay bounded. +struct Hex<'a>(&'a [u8]); + +const HEX_LIMIT: usize = 64; + +impl fmt::Debug for Hex<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for byte in self.0.iter().take(HEX_LIMIT) { + write!(f, "{byte:02x}")?; + } + if self.0.len() > HEX_LIMIT { + write!(f, "...({} bytes)", self.0.len())?; + } + Ok(()) + } +} + +/// Decodes one sparse primary key into its exact identity and `tag_columns` labels. +/// +/// This is the only decoder artifact writers use, so the identity they aggregate and +/// the labels they persist always come from the same key. +pub(crate) fn decode_series_identity( + primary_key: &[u8], + tag_columns: &[(ColumnId, String)], + offsets: &mut SparseOffsetsCache, + buf: &mut Vec, +) -> Result<(ExactSeriesKey, Vec>)> { + let mut view = SparsePrimaryKeyView::new(primary_key, offsets).context(DecodeSnafu)?; + // Validate every label once and cache its offset; the tag lookups below reuse it, + // and a truncated key never reaches a persisted row. + view.visit_label_ids(|_| Ok(())).context(DecodeSnafu)?; + let key = ExactSeriesKey::from_view(&view, primary_key); + let mut tags = Vec::with_capacity(tag_columns.len()); + for (column_id, _) in tag_columns { + tags.push( + view.label(*column_id, buf) + .context(DecodeSnafu)? + .map(str::to_owned), + ); + } + Ok((key, tags)) +} + +/// Decodes only the reserved `(table_id, tsid)` prefix, without allocating the key. +/// +/// Callers that already compared the raw bytes must not use this to decide identity; +/// it only exists where the coarse bucket is the output. +pub(crate) fn decode_series_id( + primary_key: &[u8], + offsets: &mut SparseOffsetsCache, +) -> Result { + let view = SparsePrimaryKeyView::new(primary_key, offsets).context(DecodeSnafu)?; + Ok(MetricSeriesId { + table_id: view.table_id(), + tsid: view.tsid(), + }) +} + +#[cfg(test)] +mod tests { + use store_api::codec::PrimaryKeyEncoding; + + use super::*; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn metadata() -> std::sync::Arc { + std::sync::Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )) + } + + fn key(tags: &[&str], table_id: u32, tsid: u64) -> Vec { + new_sparse_primary_key(tags, &metadata(), table_id, tsid) + } + + #[test] + fn identity_is_the_exact_key_not_the_tsid() { + let mut offsets = SparseOffsetsCache::new(); + let first = + ExactSeriesKey::from_sparse_key(&key(&["host-020078", "z"], 1, 7), &mut offsets) + .unwrap(); + let partner = + ExactSeriesKey::from_sparse_key(&key(&["host-029838", "z"], 1, 7), &mut offsets) + .unwrap(); + assert_eq!(first.id(), partner.id()); + assert_ne!(first, partner); + assert_ne!(first.bytes(), partner.bytes()); + let same = ExactSeriesKey::from_sparse_key(&key(&["host-020078", "z"], 1, 7), &mut offsets) + .unwrap(); + assert_eq!(first, same); + } + + #[test] + fn order_groups_equal_identifiers_and_sorts_groups_by_identifier() { + let mut offsets = SparseOffsetsCache::new(); + let mut keys = Vec::new(); + for (table_id, tsid, label) in [ + (1, 2, "b"), + (1, 2, "a"), + (1, 3, "a"), + (2, 1, "a"), + (1, 2, "c"), + ] { + keys.push( + ExactSeriesKey::from_sparse_key(&key(&[label, "z"], table_id, tsid), &mut offsets) + .unwrap(), + ); + } + let mut sorted = keys.clone(); + sorted.sort(); + // Byte order sorts by (table_id, tsid) first, then by the encoded labels, and + // every key sharing an identifier stays adjacent. + let ids: Vec<_> = sorted.iter().map(|key| key.id()).collect(); + assert_eq!( + ids, + vec![ + MetricSeriesId { + table_id: 1, + tsid: 2 + }, + MetricSeriesId { + table_id: 1, + tsid: 2 + }, + MetricSeriesId { + table_id: 1, + tsid: 2 + }, + MetricSeriesId { + table_id: 1, + tsid: 3 + }, + MetricSeriesId { + table_id: 2, + tsid: 1 + }, + ] + ); + // Ordering is decided on the encoded bytes alone. + for pair in keys.windows(2) { + assert_eq!(pair[0].cmp(&pair[1]), pair[0].bytes().cmp(pair[1].bytes())); + } + assert!(sorted.is_sorted()); + } + + #[test] + fn long_keys_are_identity_not_an_error() { + let mut offsets = SparseOffsetsCache::new(); + // One 8 KiB label already encodes to a sparse key past 8 KiB. Both entry points + // must accept it: the containers that store a key bound its length, and a cap on + // one entry point alone would publish keys another refuses. + let label = "x".repeat(8 * 1024); + let encoded = key(&[&label, "z"], 1, 7); + assert!(encoded.len() > 8 * 1024, "{}", encoded.len()); + let exact = ExactSeriesKey::from_sparse_key(&encoded, &mut offsets).unwrap(); + assert_eq!(encoded, exact.bytes()); + let tag_columns = [(0, "tag_0".to_string()), (1, "tag_1".to_string())]; + let (decoded, tags) = + decode_series_identity(&encoded, &tag_columns, &mut offsets, &mut Vec::new()).unwrap(); + assert_eq!(exact, decoded); + assert_eq!(Some(label), tags[0]); + } + + #[test] + fn malformed_keys_are_rejected() { + let mut offsets = SparseOffsetsCache::new(); + assert!(ExactSeriesKey::from_sparse_key(&[], &mut offsets).is_err()); + assert!(ExactSeriesKey::from_sparse_key(&[0; 13], &mut offsets).is_err()); + let good = key(&["a", "b"], 1, 1); + assert!(ExactSeriesKey::from_sparse_key(&good[..good.len() - 1], &mut offsets).is_err()); + assert_eq!( + decode_series_id(&good, &mut offsets).unwrap(), + MetricSeriesId { + table_id: 1, + tsid: 1 + } + ); + } +} diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index 777de51bfcb..f7e9781f71b 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -287,14 +287,23 @@ async fn build_index_version( .await .context(crate::error::OpenDalSnafu)?; } - // Manual builds must repair schema changes even when SST coverage is unchanged. - if manual - && reusable + // Columnar artifacts must be revalidated even when SST coverage is unchanged: + // an artifact written under another identity encoding, or with a stale tag + // schema, is structurally valid and ordered, so only opening it tells them + // apart. An unreadable artifact is rebuilt, never reused and never an error. + if reusable && selected_format == ArtifactFormat::ColumnarSeries && let Some(backend) = &formats.index { - let reader = backend.reader(store.clone(), handle.clone()).await?; - reusable = format::validate_series_schema(&reader.schema(), &version.metadata).is_ok(); + // An artifact written under another identity encoding, or with a stale tag + // schema, is structurally valid and ordered, so only opening it tells them + // apart. An unreadable artifact is rebuilt, never reused and never an error. + reusable = match backend.reader(store.clone(), handle.clone()).await { + Ok(reader) => { + format::validate_series_schema(&reader.schema(), &version.metadata).is_ok() + } + Err(_) => false, + }; } if !reusable { unusable.insert(entry.index_uuid); @@ -315,10 +324,15 @@ async fn build_index_version( .await .context(crate::error::OpenDalSnafu)?; } - if manual && reusable && backend.format() == ArtifactFormat::ColumnarSamples { - let reader = backend.reader(store.clone(), handle.clone()).await?; - reusable = - format::validate_sample_schema(&reader.schema(), &version.metadata).is_ok(); + // Same rule as series indexes: coverage, object existence and format name + // cannot prove that a samples artifact was written with exact identity. + if reusable && backend.format() == ArtifactFormat::ColumnarSamples { + reusable = match backend.reader(store.clone(), handle.clone()).await { + Ok(reader) => { + format::validate_sample_schema(&reader.schema(), &version.metadata).is_ok() + } + Err(_) => false, + }; } if reusable { reusable_samples.insert(*id); diff --git a/src/mito2/src/series_index/samples.rs b/src/mito2/src/series_index/samples.rs index 8028196513a..98cd7c76679 100644 --- a/src/mito2/src/series_index/samples.rs +++ b/src/mito2/src/series_index/samples.rs @@ -35,6 +35,7 @@ use crate::error::Result; use crate::series_index::artifact::{ArtifactFormat, SamplesFileHandle}; use crate::series_index::columnar::{invalid, require, unsupported}; use crate::series_index::format::{FormatWriter, SamplesFormat, SamplesFormatReader}; +use crate::series_index::identity::IdentityEncoding; use crate::series_index::{MetricSeriesId, TABLE_ID_COLUMN, TSID_COLUMN}; const VERSION: u32 = 1; @@ -44,6 +45,8 @@ const MAX_FIELDS: usize = 1024; const BLOCK_BYTES: usize = 256 * 1024; const SEQUENCE: &str = "__sequence"; const OPERATION: &str = "__op_type"; +/// Output column naming the exact series key that owns each returned sample row. +const PRIMARY_KEY: &str = store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; /// Internal packing knobs; the backend is disabled unless selected explicitly. #[derive(Debug, Clone)] @@ -151,6 +154,8 @@ struct Part { struct Manifest { format: ArtifactFormat, version: u32, + /// Mandatory discriminator; see [`IdentityEncoding`]. Never defaulted on read. + identity_encoding: IdentityEncoding, region_id: u64, source_file_id: FileId, generation: FileId, @@ -185,7 +190,13 @@ impl Manifest { (1..=MAX_FIELDS).contains(&self.fields.len()), "invalid samples field count", )?; - let mut names = HashSet::from([TABLE_ID_COLUMN, TSID_COLUMN, SEQUENCE, OPERATION]); + let mut names = HashSet::from([ + TABLE_ID_COLUMN, + TSID_COLUMN, + SEQUENCE, + OPERATION, + PRIMARY_KEY, + ]); require( names.insert(&self.timestamp.name), "duplicate samples timestamp name", @@ -221,8 +232,17 @@ impl Manifest { fields.extend(self.fields.iter().map(|f| { crate::sst::with_field_id(Field::new(&f.name, DataType::Float64, true), f.column_id) })); + // Appended after the data fields so stored row indices stay stable. The column is the + // artifact's exact-identity marker: readers attribute every row of a batch by the key + // they return with it and never project this column. + fields.push(Field::new(PRIMARY_KEY, DataType::Binary, false)); Arc::new(Schema::new(fields)) } + + /// Index of the exact series key column in [`Self::schema`]. + fn primary_key_index(&self) -> usize { + 5 + self.fields.len() + } } fn valid_rows(rows: usize, fields: usize) -> bool { diff --git a/src/mito2/src/series_index/samples/metadata.rs b/src/mito2/src/series_index/samples/metadata.rs index aaeeac3f3e9..8958de7b51b 100644 --- a/src/mito2/src/series_index/samples/metadata.rs +++ b/src/mito2/src/series_index/samples/metadata.rs @@ -22,10 +22,27 @@ use crate::series_index::samples::{ pub(super) const MAGIC: &[u8; 8] = b"GTSAMPLE"; pub(super) const HEADER_SIZE: u64 = 16; -pub(super) const TRAILER_SIZE: u64 = 112; +pub(super) const TRAILER_SIZE: u64 = 120; pub(super) const PAGE_LIMIT: usize = 65536; pub(super) const ROOT_CAPACITY: usize = (PAGE_LIMIT - 12) / 72; const CHUNK_LIMIT: usize = 1024 * 1024; +/// Fixed part of a block descriptor, before its 44-byte column descriptors. +/// +/// The exact series key is not part of it. A descriptor keeps a fixed-size [`KeyRef`] into +/// the part's key area instead, so descriptor size is independent of key length and a key +/// longer than a metadata page is representable. +pub(super) const BLOCK_FIXED_BYTES: usize = 56; +/// Smallest possible block descriptor: the fixed part plus its three mandatory columns. +const MIN_BLOCK_BYTES: usize = BLOCK_FIXED_BYTES + 3 * 44; + +/// Descriptor bytes for a block with `columns` column descriptors. +/// +/// Independent of the key's length: the descriptor references the key, it does not hold it. +/// This is the one place the descriptor size formula lives, for the writer's page planning +/// and for readers that inspect decoded pages. +pub(super) fn descriptor_bytes(columns: usize) -> usize { + BLOCK_FIXED_BYTES + 44 * columns +} #[derive(Debug, Clone, Copy)] pub(super) struct PayloadReference { @@ -201,32 +218,219 @@ pub(super) struct Column { pub payload: PayloadReference, pub nulls: u32, } +/// The location of one exact series key inside a part's key area. +/// +/// This reference is the descriptor's whole identity. The key bytes themselves are stored +/// once per series in the part's key area, outside every bounded metadata page, so a key may +/// be longer than a page and the descriptor stays 12 bytes whatever its length is. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct KeyRef { + /// Absolute offset of the first key byte from the start of the part. + pub offset: u64, + /// Key length in bytes; nonzero. + pub length: u32, +} + +impl KeyRef { + pub fn end(&self) -> Result { + self.offset + .checked_add(u64::from(self.length)) + .ok_or_else(|| invalid("sample key offset overflow")) + } + + /// End of this key's stored entry: its bytes followed by its checksum. + /// + /// `length` covers the key only, so the stored checksum is exactly [`KEY_CHECKSUM_BYTES`] + /// past [`Self::end`] and never part of what the reference itself names. + pub fn checksum_end(&self) -> Result { + self.end()? + .checked_add(KEY_CHECKSUM_BYTES) + .ok_or_else(|| invalid("sample key offset overflow")) + } + + fn encode(&self, out: &mut Vec) { + put_u64(out, self.offset); + put_u32(out, self.length); + } + + fn decode(c: &mut Cursor<'_>) -> Result { + Ok(Self { + offset: c.u64()?, + length: c.u32()?, + }) + } + + /// Requires this reference to name a nonempty key entry inside `area`. + /// + /// The entry includes the stored checksum, so a reference whose key ends inside the area but + /// whose checksum would not is rejected like any other out-of-area reference. + pub fn validate(&self, area: KeyArea) -> Result<()> { + require( + self.length > 0 && self.offset >= area.start && self.checksum_end()? <= area.end, + "sample key reference outside key area", + ) + } +} + +/// Bytes of the XXH64 stored after each key, seed zero. +/// +/// A key copy is `[key bytes][checksum]`; the descriptor's reference names only the key, so a +/// fetched key can be verified before it is used to attribute rows to a series. Without it a +/// flipped label byte would silently move rows to another series, which the old inline-key layout +/// could not suffer: its keys sat in a page whose stored bytes were already checksummed. +pub(super) const KEY_CHECKSUM_BYTES: u64 = std::mem::size_of::() as u64; + +/// A part's key area: exact series keys stored back to back, outside every metadata page. +/// +/// Keys are appended in block order and are never page-bounded, so their length is bounded by +/// the part, not by `PAGE_LIMIT`. Every copy is followed by its checksum. Because blocks are +/// ordered by exact key, every block of one series is adjacent and shares a single copy of the +/// key, which makes an equal [`KeyRef`] mean "the same exact series". +#[derive(Debug, Clone, Copy)] +pub(super) struct KeyArea { + pub start: u64, + pub end: u64, +} + +/// The three regions of one part, in file order: column payloads, exact series keys, pages. +#[derive(Debug, Clone, Copy)] +pub(super) struct Areas { + /// End of the column payload area and start of the key area. + pub payload_end: u64, + /// The key area every block descriptor references. + pub keys: KeyArea, + /// Start of the metadata area, which is the end of the key area. + pub metadata_start: u64, +} + #[derive(Debug, Clone)] pub(super) struct Block { + /// Coarse `(table_id, tsid)` prefix decoded from the block's exact key. + /// + /// It prunes candidates but never names the owner: `(table_id, tsid)` is a hash, so two + /// distinct series can share it. The owner is the key [`Self::key`] references. pub id: MetricSeriesId, + /// Fixed-size reference to the block's exact key in the part's key area. + /// + /// Equal references mean the same exact series; a different series that collides on the + /// identifier always carries a different reference. + pub key: KeyRef, pub min: i64, pub max: i64, pub ordinal: u64, pub rows: u32, pub columns: Vec, } -/// The ordering boundary retained between pages, without their column descriptors. -#[derive(Clone, Copy)] +/// The ordering boundary retained between pages and parts, without column descriptors. +/// +/// A [`KeyRef`] is an offset inside one part's key area, so two references are comparable only +/// when they belong to the same part; `part` carries that scope in the boundary itself so no +/// caller can compare references across a part boundary by accident. Two parts hold unrelated +/// key areas, which is why a series that spans parts is recognised by [`BlockOrder::continues`] +/// on the exact key bytes instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct BlockOrder { + part: u32, id: MetricSeriesId, + key: KeyRef, + min: i64, max: i64, ordinal: u64, } +impl BlockOrder { + /// The part this block belongs to. + pub(crate) fn part(&self) -> u32 { + self.part + } + + /// The coarse `(table_id, tsid)` prefix, which is ordered across parts as well. + pub(crate) fn id(&self) -> MetricSeriesId { + self.id + } + + /// The block's key reference, a location inside [`Self::part`]'s key area. + pub(crate) fn key(&self) -> KeyRef { + self.key + } + + /// Checks the next block of the *same* part against this boundary. + /// + /// Ownership and ordering of a part are decided on the exact `(table_id, tsid)` and on the + /// key reference only: two blocks of one part that share one reference share one key copy, + /// so they are one exact series, while a different series always has its own copy, placed + /// after the previous one and restarting the ordinal. Only descriptor fields are read, so a + /// block this request prunes is never fetched and its key content is not validated. + pub(crate) fn precedes(&self, next: &Self, adjacent: bool) -> Result<()> { + require( + self.part == next.part, + "sample key references compared across parts", + )?; + require(self.id <= next.id, "unordered sample series keys")?; + if next.key == self.key { + self.continues(next, adjacent)?; + } else { + // A different series must place its one key copy after the previous series' copy. + // Overlap, or a second copy of a key that is still current, is a corrupt reference + // rather than a silent second read of the same bytes. + require( + self.key.checksum_end()? <= next.key.offset, + "overlapping sample key references", + )?; + if adjacent { + require(next.ordinal == 0, "first sample block has nonzero ordinal")?; + } + } + Ok(()) + } + + /// Requires `next` to continue one exact series in time and ordinal. + /// + /// Equal exact key bytes are what make two blocks the same series across a part boundary, + /// where their references belong to different key areas. A repeated key that restarts its + /// ordinal or goes back in time is a corrupt part, not a second series. + pub(crate) fn continues(&self, next: &Self, adjacent: bool) -> Result<()> { + require( + self.max <= next.min && self.ordinal < next.ordinal, + "unordered sample blocks", + )?; + if adjacent { + require( + self.ordinal.checked_add(1) == Some(next.ordinal), + "nonconsecutive sample ordinal", + )?; + } + Ok(()) + } +} + impl Block { - pub(crate) fn order(&self) -> BlockOrder { + /// The ordering boundary this block contributes, scoped to its part. + pub(crate) fn order(&self, part: u32) -> BlockOrder { BlockOrder { + part, id: self.id, + key: self.key, + min: self.min, max: self.max, ordinal: self.ordinal, } } + /// Checks this block against the previous block of the same page. + /// + /// A page never spans parts, so both references name one key area and are comparable; a block + /// in another part is compared through [`BlockOrder::continues`] and its exact key instead. + fn follows_in_page(&self, previous: &Self, adjacent: bool) -> Result<()> { + // One scope for both sides: a page belongs to exactly one part, and no reference built + // here outlives the page. + const PAGE_SCOPE: u32 = 0; + previous + .order(PAGE_SCOPE) + .precedes(&self.order(PAGE_SCOPE), adjacent) + } + + /// The coarse bounds used for pruning. Exact ownership lives in `key`. pub fn bounds(&self) -> Bounds { Bounds { first: self.id, @@ -235,12 +439,10 @@ impl Block { max: self.max, } } - pub fn size(&self) -> usize { - 44 + 44 * self.columns.len() - } pub fn encode(&self, out: &mut Vec) { put_u32(out, self.id.table_id); put_u64(out, self.id.tsid); + self.key.encode(out); put_u64(out, self.min as u64); put_u64(out, self.max as u64); put_u64(out, self.ordinal); @@ -256,13 +458,19 @@ impl Block { fn decode( c: &mut Cursor<'_>, manifest: &Manifest, - payload_end: u64, + areas: Areas, previous_end: &mut u64, ) -> Result { let id = MetricSeriesId { table_id: c.u32()?, tsid: c.u64()?, }; + // Only the key's location is decoded here. The bytes stay in the key area, so a + // descriptor never grows with the key and a request that prunes this block never + // reads it; the reader validates the reference and its `(table_id, tsid)` prefix + // when it fetches the key for a block it actually selects. + let key = KeyRef::decode(c)?; + key.validate(areas.keys)?; let min = c.u64()? as i64; let max = c.u64()? as i64; let ordinal = c.u64()?; @@ -317,7 +525,7 @@ impl Block { }; require(valid, "sample codec role mismatch")?; } - payload.validate(*previous_end, payload_end, false)?; + payload.validate(*previous_end, areas.payload_end, false)?; *previous_end = payload.end()?; columns.push(Column { role, @@ -328,6 +536,7 @@ impl Block { } Ok(Self { id, + key, min, max, ordinal, @@ -335,24 +544,6 @@ impl Block { columns, }) } - pub fn follows(&self, previous: &BlockOrder, adjacent: bool) -> Result<()> { - require(previous.id <= self.id, "unordered sample identifiers")?; - if self.id == previous.id { - require( - previous.max <= self.min && previous.ordinal < self.ordinal, - "unordered sample blocks", - )?; - if adjacent { - require( - previous.ordinal.checked_add(1) == Some(self.ordinal), - "nonconsecutive sample ordinal", - )?; - } - } else if adjacent { - require(self.ordinal == 0, "first sample block has nonzero ordinal")?; - } - Ok(()) - } } #[derive(Debug)] pub(super) enum Page { @@ -368,6 +559,8 @@ impl Page { } Self::Leaf(blocks) => { blocks.capacity() * std::mem::size_of::() + // Each block owns its column descriptors on the heap. Its exact + // series key is not retained: it stays in the part's key area. + blocks .iter() .map(|block| block.columns.capacity() * std::mem::size_of::()) @@ -407,11 +600,16 @@ impl Page { } out } + /// Decodes one metadata page. + /// + /// `areas` bounds what the page may reference: column payloads end where the key area + /// starts, keys live in their own area, and the page area starts at `areas.metadata_start` + /// so every child page must lie after it and before its parent. pub fn decode( bytes: &[u8], reference: &PageReference, depth: u32, - payload_end: u64, + areas: Areas, manifest: &Manifest, ) -> Result { require(bytes.len() <= PAGE_LIMIT, "sample page exceeds limit")?; @@ -422,15 +620,15 @@ impl Page { let page = match (kind, depth) { (0, 1) => { require( - count <= c.remaining() / 176, + count <= c.remaining() / MIN_BLOCK_BYTES, "sample block count exceeds page", )?; let mut blocks: Vec = Vec::with_capacity(count); let mut end = HEADER_SIZE; for _ in 0..count { - let block = Block::decode(&mut c, manifest, payload_end, &mut end)?; + let block = Block::decode(&mut c, manifest, areas, &mut end)?; if let Some(previous) = blocks.last() { - block.follows(&previous.order(), true)?; + block.follows_in_page(previous, true)?; } blocks.push(block); } @@ -442,7 +640,7 @@ impl Page { "sample root exceeds capacity", )?; let mut children: Vec = Vec::with_capacity(count); - let mut end = payload_end; + let mut end = areas.metadata_start; for _ in 0..count { let child = PageReference::decode(&mut c)?; child @@ -488,7 +686,10 @@ pub(super) struct Trailer { pub root: PageReference, pub depth: u32, pub file_length: u64, + /// End of the payload area, which is where the key area starts. pub payload_end: u64, + /// End of the key area, which is where the metadata pages start. + pub key_end: u64, } impl Trailer { pub fn encode(&self) -> Vec { @@ -498,6 +699,7 @@ impl Trailer { put_u32(&mut out, VERSION); put_u64(&mut out, self.file_length); put_u64(&mut out, self.payload_end); + put_u64(&mut out, self.key_end); out.extend(MAGIC); let sum = codec::checksum(&out); put_u64(&mut out, sum); @@ -517,18 +719,26 @@ impl Trailer { } let length = c.u64()?; let payload_end = c.u64()?; + let key_end = c.u64()?; require( - c.take(8)? == MAGIC && c.u64()? == codec::checksum(&bytes[..104]), + // The checksum covers everything before it, magic included: the trailer body plus + // the 8-byte magic, which is `TRAILER_SIZE - 8` and not a frozen constant. + c.take(8)? == MAGIC + && c.u64()? + == codec::checksum( + &bytes[..TRAILER_SIZE as usize - std::mem::size_of::()], + ), "sample trailer checksum/magic", )?; require( length == file_length && payload_end >= HEADER_SIZE + && key_end >= payload_end && file_length >= HEADER_SIZE + TRAILER_SIZE, "invalid sample framing", )?; root.payload - .validate(payload_end, file_length - TRAILER_SIZE, true)?; + .validate(key_end, file_length - TRAILER_SIZE, true)?; require( root.payload.end()? == file_length - TRAILER_SIZE, "sample root must end at trailer", @@ -538,6 +748,24 @@ impl Trailer { depth, file_length, payload_end, + key_end, }) } + + /// The key area every block descriptor in this part references. + pub fn keys(&self) -> KeyArea { + KeyArea { + start: self.payload_end, + end: self.key_end, + } + } + + /// The part's three regions, as page decoding needs them. + pub fn areas(&self) -> Areas { + Areas { + payload_end: self.payload_end, + keys: self.keys(), + metadata_start: self.key_end, + } + } } diff --git a/src/mito2/src/series_index/samples/reader.rs b/src/mito2/src/series_index/samples/reader.rs index e5be1f0c99c..27efd0a7648 100644 --- a/src/mito2/src/series_index/samples/reader.rs +++ b/src/mito2/src/series_index/samples/reader.rs @@ -14,13 +14,13 @@ //! Lazy, selective sample reads with independent traversal and pinned lifetimes. -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::sync::Arc; use std::time::Instant; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datatypes::arrow::array::{ - ArrayRef, Float64Array, TimestampMicrosecondArray, TimestampMillisecondArray, + ArrayRef, BinaryArray, Float64Array, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt32Array, UInt64Array, }; use datatypes::arrow::datatypes::SchemaRef; @@ -33,13 +33,14 @@ use tokio::sync::OnceCell; use crate::cache::columnar::{ self as metadata_cache, MetadataCache, PayloadCache, PayloadPartCache, }; -use crate::error::{NewRecordBatchSnafu, OpenDalSnafu, Result, SerdeJsonSnafu}; +use crate::error::{NewRecordBatchSnafu, OpenDalSnafu, Result}; use crate::metrics::ColumnarReadMetrics; use crate::series_index::artifact::{ArtifactFormat, SamplesFileHandle}; -use crate::series_index::format::{SamplesFormatReader, SamplesReadRequest}; +use crate::series_index::format::{SamplesBatch, SamplesFormatReader, SamplesReadRequest}; +use crate::series_index::identity::{ExactSeriesKey, IdentityEncoding}; use crate::series_index::samples::codec; use crate::series_index::samples::metadata::{ - self, Block, BlockOrder, Bounds, Column, Page, PageReference, Trailer, + self, Areas, Block, BlockOrder, Bounds, Column, Page, PageReference, Trailer, }; use crate::series_index::samples::{ MANIFEST_LIMIT, Manifest, MetricSeriesId, TimestampUnit, VERSION, invalid, require, unsupported, @@ -134,10 +135,23 @@ impl SamplesReader { "samples manifest exceeds limit", )?; let bytes = read_range(&store, &path, 0, length).await?; - let manifest: Manifest = serde_json::from_slice(&bytes).context(SerdeJsonSnafu)?; + // A manifest that the current format cannot read predates exact identity or is + // corrupt; either way it must never be reused. + let manifest: Manifest = serde_json::from_slice(&bytes).map_err(|error| { + invalid(format!( + "samples manifest is not readable; rebuild the artifact: {error}" + )) + })?; if manifest.version != VERSION { return Err(unsupported("samples manifest version")); } + // Old artifacts can be structurally valid and perfectly ordered, so identity + // encoding is the only thing that keeps their lost labels out of reads. + if manifest.identity_encoding != IdentityEncoding::CURRENT_SAMPLES { + return Err(unsupported( + "samples manifest identity encoding; rebuild the samples artifact", + )); + } require( manifest.format == artifact.format && manifest.region_id == region @@ -195,7 +209,10 @@ impl SamplesFormatReader for SamplesReader { fn schema(&self) -> SchemaRef { self.schema.clone() } - fn read(&self, request: SamplesReadRequest) -> Result>> { + fn read( + &self, + request: SamplesReadRequest, + ) -> Result>> { let schema = self.schema(); let mut projection = Vec::new(); let mut seen = HashSet::new(); @@ -207,6 +224,10 @@ impl SamplesFormatReader for SamplesReader { .map_err(|_| invalid(format!("unknown sample projection {name}")))?, ); } + require( + !projection.contains(&self.manifest.primary_key_index()), + "the exact series key is returned once per batch, not projected as a column", + )?; require( request .start @@ -214,12 +235,18 @@ impl SamplesFormatReader for SamplesReader { .is_none_or(|(start, end)| start <= end), "reversed sample time range", )?; - let empty = request.ids.as_ref().is_some_and(|ids| ids.is_empty()) + let empty = request.keys.as_ref().is_some_and(|keys| keys.is_empty()) || request.start.zip(request.end).is_some_and(|(a, b)| a == b); - let complete = request.ids.is_none() + let complete = request.keys.is_none() && request.table_id.is_none() && request.start.is_none() && request.end.is_none(); + // Coarse pruning bounds come from the exact selection itself, so a reader can never + // prune away a part or page that owns one of the selected series. + let ids: Option> = request + .keys + .as_ref() + .map(|keys| keys.iter().map(ExactSeriesKey::id).collect()); let projected = Arc::new( schema .project(&projection) @@ -228,7 +255,10 @@ impl SamplesFormatReader for SamplesReader { let field_ids: HashSet<_> = projection .iter() .filter(|&&index| index >= 5) - .map(|&index| self.manifest.fields[index - 5].column_id) + // The exact series key column is appended after the data fields and has no + // column id of its own. + .filter_map(|&index| self.manifest.fields.get(index - 5)) + .map(|field| field.column_id) .collect(); let cache = self.cache.clone(); let manifest = self.manifest.clone(); @@ -240,11 +270,26 @@ impl SamplesFormatReader for SamplesReader { let _handle = &handle; let mut metrics = ColumnarReadMetrics::default(); if empty { return; } + // Traversal state: the last block this stream visited, whatever it selected, and + // whether it is the block immediately before the next one in the artifact. let mut previous: Option = None; let mut adjacent = true; + // Content state: the last block whose exact key this stream loaded. Selection can + // skip a block inside a page, a whole page, or a whole part, so this state is not + // advanced by a block whose key was pruned: the two loaded keys on either side of a + // gap are still compared with each other. + let mut loaded_key: Option = None; + // The last loaded key of the previous part, with the pool bytes it is charged. A + // reference is a location inside one part, so the exact bytes are what decides + // whether a series continues into another; this copy outlives every per-part + // reservation, so it has a consumer of its own. + let boundary_reservation = request.memory_pool.as_ref() + .map(|pool| MemoryConsumer::new("ColumnarSamplesBoundaryKey").register(pool)); + let mut spanning_key: Option = None; + let mut spanning_bytes = 0usize; for part in &manifest.files { let bounds = Bounds { first: part.first, last: part.last, min: part.min_timestamp, max: part.max_timestamp }; - if !matches(bounds, &request) { adjacent = false; continue; } + if !matches(bounds, &request, ids.as_ref()) { adjacent = false; continue; } let path = manifest.part_path(part.part_id); // One lazy reader per accessed part, shared across this query's streams. let reader = parts[part.part_id as usize].get_or_try_init(|| async { @@ -263,57 +308,486 @@ impl SamplesFormatReader for SamplesReader { reader, cache: payload_cache.as_ref().map(|cache| cache.part(handle.entry().artifact.generation, part.part_id, size)), end: trailer.payload_end, + key_end: trailer.key_end, }; require(trailer.root.bounds == bounds, "sample root/part bounds mismatch")?; - let root = load_metadata_page(cache.as_ref(), reader, (&path, size), &trailer.root, trailer.depth, trailer.payload_end, &manifest).await?; + let root = load_metadata_page(cache.as_ref(), reader, (&path, size), &trailer.root, trailer.depth, trailer.areas(), &manifest).await?; metrics.add_elapsed("columnar_metadata_wall", metadata_started.elapsed()); let (mut active, children) = match root.as_ref() { Page::Leaf(_) => (Some(root.clone()), Vec::new()), Page::Internal(children) => (None, children.clone()) }; let mut children = children.into_iter(); let mut row_count = 0u64; let mut block_count = 0u64; let mut payload_position = metadata::HEADER_SIZE; + // Exact keys live in the part's key area, outside every metadata page. Every + // block of one series references the same copy, so one load serves them all. + let keys_reservation = request.memory_pool.as_ref() + .map(|pool| MemoryConsumer::new("ColumnarSamplesKeys").register(pool)); + // `KeyRef` is the cache key; its fields are readable, so the map needs no + // ordering or hashing of the reference type itself. + let mut part_keys = HashMap::<(u64, u32), ExactSeriesKey>::new(); + let mut key_offsets = mito_codec::row_converter::SparseOffsetsCache::new(); loop { let page = if let Some(page) = active.take() { page } else { let Some(child) = children.next() else { break; }; - if !matches(child.bounds, &request) { adjacent = false; continue; } + if !matches(child.bounds, &request, ids.as_ref()) { adjacent = false; continue; } let started = Instant::now(); - let page = load_metadata_page(cache.as_ref(), reader, (&path, size), &child, 1, trailer.payload_end, &manifest).await?; + let page = load_metadata_page(cache.as_ref(), reader, (&path, size), &child, 1, trailer.areas(), &manifest).await?; metrics.add_elapsed("columnar_metadata_wall", started.elapsed()); page }; let blocks = match page.as_ref() { Page::Leaf(blocks) => blocks, Page::Internal(_) => Err(unsupported("internal child in sample depth-two tree"))? }; + // One merged read for the keys of this page's selected blocks: a page can + // hold hundreds of short series, and one read per series would multiply + // this scan's I/O by the number of series it touches. + load_page_keys( + &payload_reader, + trailer.areas(), + blocks, + &request, + ids.as_ref(), + &mut part_keys, + keys_reservation.as_ref(), + &mut key_offsets, + &mut metrics, + ) + .await?; let mut prefetched = PrefetchedBlocks::default(); for (block_index, block) in blocks.iter().enumerate() { row_count = row_count.checked_add(u64::from(block.rows)).ok_or_else(|| invalid("samples row count overflow"))?; block_count += 1; - if let Some(prior) = &previous { block.follows(prior, adjacent)?; } - else if complete { require(block.ordinal == 0, "initial sample ordinal is not zero")?; } + let order = block.order(part.part_id); + if let Some(prior) = &previous { + if prior.part() == order.part() { + // Reference positions are part-relative: only a block of the same + // part may be compared with the previous block's reference. + prior.precedes(&order, adjacent)?; + } else { + // Across a part boundary the references name unrelated key areas, + // so only the ordered identifier applies here. Whether the same + // series continues is decided on the exact keys below. + require(prior.id() <= order.id(), "unordered sample series keys")?; + } + } else if complete { require(block.ordinal == 0, "initial sample ordinal is not zero")?; } for column in &block.columns { column.payload.validate(payload_position, trailer.payload_end, false)?; payload_position = column.payload.end()?; } - let batch = if matches(block.bounds(), &request) { + // Ownership is exact, so it needs the block's key rather than its + // coarse identifier. Pruned blocks never load one, and every key this + // page selects was already fetched in one merged read. + let key_id = (block.key.offset, block.key.length); + let mut owner = false; + if matches(block.bounds(), &request, ids.as_ref()) { + if let std::collections::hash_map::Entry::Vacant(slot) = + part_keys.entry(key_id) + { + // The page-level prefetch was refused by the memory pool, or + // this series has no copy cached for the part yet. This path + // reserves its own read before issuing it. + let key = load_block_key( + &payload_reader, + trailer.areas(), + block, + keys_reservation.as_ref(), + &mut key_offsets, + &mut metrics, + ) + .await?; + slot.insert(key); + } + let key = part_keys + .get(&key_id) + .ok_or_else(|| invalid("missing sample key"))?; + // A key that parses but claims another identifier is a malformed + // reference: ownership would silently point at the wrong series. + require( + key.id() == block.id, + "sample block series key disagrees with its identifier", + )?; + // Only the keys this stream loaded can be checked here: a reference + // is a location, and selection prunes blocks without reading theirs, + // so unvisited ranges are validated structurally only. The check is + // chained across such gaps - it compares the last *loaded* key, not + // the last visited block - so skipping a block cannot hide the keys + // on either side of it. Loaded keys must be nondecreasing, and a + // repeated key must continue its series instead of restarting it. + // That is also the only way a series that spans a part boundary is + // recognised: its two references name unrelated key areas. + if let Some(prior) = &loaded_key { + let prior_key = if prior.part() == order.part() { + let reference = prior.key(); + part_keys.get(&(reference.offset, reference.length)) + } else { + spanning_key.as_ref() + }; + if let Some(prior_key) = prior_key { + // The exact-next-block rules only hold when nothing was + // traversed in between; a gap still requires the order and + // the ordinal to keep increasing. + let consecutive = + adjacent && previous.as_ref() == Some(prior); + require(key >= prior_key, "unordered sample exact series keys")?; + if key == prior_key { + prior.continues(&order, consecutive)?; + } else if consecutive { + // A new series starts at ordinal zero, which the + // structural check above only covers inside one part. + require(block.ordinal == 0, "first sample block has nonzero ordinal")?; + } + } + } + loaded_key = Some(order); + owner = owns(key.bytes(), &request); + } + let batch = if owner { if block_index >= prefetched.end { // Release the previous window before reserving the next one. drop(std::mem::take(&mut prefetched)); - prefetched = prefetch_blocks(&payload_reader, blocks, block_index, &field_ids, &request, &projection, trailer.payload_end, &mut metrics).await?; + prefetched = prefetch_blocks(&payload_reader, blocks, block_index, &field_ids, &request, ids.as_ref(), &projection, trailer.payload_end, &mut metrics).await?; } *metrics.counts.entry("blocks").or_default() += 1; *metrics.counts.entry("decoded_rows").or_default() += u64::from(block.rows); read_block(&payload_reader, block, &manifest, &request, &projection, &field_ids, projected.clone(), &mut prefetched, &mut metrics).await? } else { None }; if let Some((batch, _reservation)) = batch { + // The key is copied only now: `read_block` has just reserved the + // bytes the batch retains, so the clone is paid for before it + // happens instead of being allocated first and charged after. + let key = part_keys + .get(&key_id) + .ok_or_else(|| invalid("missing sample key"))? + .clone(); *metrics.counts.entry("batches").or_default() += 1; *metrics.counts.entry("output_rows").or_default() += batch.num_rows() as u64; - yield batch; + // The block's key is emitted once, not once per row. + yield SamplesBatch { batch, key }; } - previous = Some(block.order()); adjacent = true; + previous = Some(order); adjacent = true; } } + // The next part's references name its own key area, so the last loaded block's + // exact key is the only thing that can prove a series continues there. It is + // kept, with its own pool charge, only while another part remains to be read. A + // part that loaded nothing leaves the carried key alone: it is still the last + // key this stream read. + if let Some(loaded) = loaded_key.filter(|loaded| loaded.part() == part.part_id) { + let retain = ((part.part_id as usize) + 1 < manifest.files.len()) + .then(|| { + let reference = loaded.key(); + part_keys.get(&(reference.offset, reference.length)) + }) + .flatten(); + retain_boundary_key( + boundary_reservation.as_ref(), + &mut spanning_bytes, + &mut spanning_key, + retain, + )?; + } if complete { require(row_count == part.num_rows && block_count == part.num_blocks, "sample part totals mismatch")?; } } })) } } -fn matches(bounds: Bounds, request: &SamplesReadRequest) -> bool { +/// Every block records the exact series key that owns its rows. A coarse identifier +/// match is not ownership: a colliding partner shares the identifier. +fn owns(key: &[u8], request: &SamplesReadRequest) -> bool { + request.keys.as_ref().is_none_or(|keys| keys.contains(key)) +} + +/// Retained bytes per cached key beyond the key bytes themselves: the boxed key and the +/// allocator's rounding of the key's own allocation. The map entry is not charged, matching the +/// per-block identity allowance, so a page of short keys still fits a small scan pool. +const KEY_ENTRY_BYTES: usize = std::mem::size_of::() + BUFFER_ALIGNMENT; + +/// Bytes of the checksum stored after each key in the key area. +const KEY_CHECKSUM_BYTES: usize = metadata::KEY_CHECKSUM_BYTES as usize; + +/// Retained bytes of one cached key: its bytes, the boxed key, and the allocator's rounding. +fn retained_key_bytes(length: u32) -> usize { + length as usize + KEY_ENTRY_BYTES +} + +/// Replaces the exact key that carries one series across a part boundary. +/// +/// Unlike the keys of a part's blocks, this copy outlives every per-part reservation, so it has a +/// consumer of its own: a replacement is charged before it is made, and the previous charge goes +/// back as soon as it is replaced or no longer needed. +fn retain_boundary_key( + reservation: Option<&MemoryReservation>, + charged: &mut usize, + slot: &mut Option, + key: Option<&ExactSeriesKey>, +) -> Result<()> { + *slot = None; + if let Some(reservation) = reservation { + let _ = reservation.try_shrink(*charged); + } + *charged = 0; + if let Some(key) = key { + let bytes = retained_key_bytes(key.bytes().len() as u32); + reserve_keys(reservation, bytes)?; + *slot = Some(key.clone()); + *charged = bytes; + } + Ok(()) +} + +/// Reserves `bytes` for a key fetch, or reports scan-memory exhaustion. +/// +/// The reservation is taken before the read, so a pool that cannot hold the fetched key entry and +/// the owned key it becomes fails without issuing I/O. +fn reserve_keys(reservation: Option<&MemoryReservation>, bytes: usize) -> Result<()> { + if let Some(reservation) = reservation { + reservation + .try_grow(bytes) + .context(crate::error::ScanMemoryExhaustedSnafu)?; + } + Ok(()) +} + +/// Reserves `bytes` for a key fetch if the pool can pay for it. +/// +/// A prefetch is optional: refusing it leaves the keys to [`load_block_key`], which is where an +/// unpayable key reports exhaustion. +fn try_reserve_keys(reservation: Option<&MemoryReservation>, bytes: usize) -> bool { + reservation.is_none_or(|reservation| reservation.try_grow(bytes).is_ok()) +} + +/// Releases the transient part of a key-fetch reservation. +/// +/// The fetched bytes and the owned keys decoded from them coexist, so the pool is charged the +/// peak up front. Once the buffers are dropped only the owned keys remain, and the difference +/// goes back to the pool. +struct KeyPeak<'a> { + reservation: Option<&'a MemoryReservation>, + transient: usize, +} +impl Drop for KeyPeak<'_> { + fn drop(&mut self) { + if let Some(reservation) = self.reservation + && self.transient > 0 + { + let _ = reservation.try_shrink(self.transient); + } + } +} + +/// Splits one stored key entry into its key bytes after verifying the stored checksum. +/// +/// The reference's length covers the key only, so the entry is exactly one key followed by +/// [`KEY_CHECKSUM_BYTES`]. A flipped label byte keeps the key's framing and identifier valid, so +/// this is the only place that notices it before the bytes attribute rows to a series. +fn verified_key(entry: &[u8], length: u32) -> Result<&[u8]> { + let length = length as usize; + require( + entry.len() == length + KEY_CHECKSUM_BYTES, + "sample key entry length mismatch", + )?; + let (key, stored) = entry.split_at(length); + let stored: [u8; KEY_CHECKSUM_BYTES] = stored + .try_into() + .map_err(|_| invalid("sample key checksum framing"))?; + require( + codec::checksum(key) == u64::from_le_bytes(stored), + "sample key checksum mismatch", + )?; + Ok(key) +} + +/// Reads one block's exact key from the part's key area. +/// +/// The page-level prefetch serves a whole page in one read, but it is charged up front. When the +/// scan pool cannot take that, this path reads one key at a time, so a block that fits the pool +/// stays readable. Either way the pool is charged the same coexisting peak before any I/O: the +/// fetched entry, whatever alignment or assembly that fetch costs, and the owned key built from +/// it. Reading first and refusing the key afterwards would leave the bytes already fetched. +async fn load_block_key( + reader: &PayloadReader<'_>, + areas: Areas, + block: &Block, + reservation: Option<&MemoryReservation>, + offsets: &mut mito_codec::row_converter::SparseOffsetsCache, + metrics: &mut ColumnarReadMetrics, +) -> Result { + block.key.validate(areas.keys)?; + let retained = retained_key_bytes(block.key.length); + let range = block.key.offset..block.key.checksum_end()?; + let ranges = std::slice::from_ref(&range); + // The aligned cache turns one key into a whole aligned block plus whatever assembling it + // needs, so a warm part is asked for that peak first. A pool that cannot cover it gets an + // explicitly cache-less reader instead, whose peak is the key entry alone; only then does the + // read run. + let mut aligned = reader.cache.is_some(); + let mut peak = fetch_peak(reader, ranges, areas.keys.end, aligned) + retained; + if !try_reserve_keys(reservation, peak) { + // Alignment is optional. Retrying on the smaller budget also reports exhaustion when even + // one key entry does not fit. + aligned = false; + peak = fetch_peak(reader, ranges, areas.keys.end, false) + retained; + reserve_keys(reservation, peak)?; + } + let _peak = KeyPeak { + reservation, + transient: peak - retained, + }; + let uncached = PayloadReader { + reader: reader.reader, + cache: None, + end: reader.end, + key_end: reader.key_end, + }; + let target = if aligned { reader } else { &uncached }; + let mut fetched = fetch_ranges(target, areas.keys.end, ranges, metrics).await?; + let entry = fetched + .pop() + .ok_or_else(|| invalid("missing sample key response"))?; + let key = ExactSeriesKey::from_sparse_key(verified_key(&entry, block.key.length)?, offsets)?; + // A key that parses but claims another identifier is a malformed reference: ownership would + // silently point at the wrong series. + require( + key.id() == block.id, + "sample block series key disagrees with its identifier", + )?; + Ok(key) +} + +/// Peak bytes [`fetch_ranges`] can hold at once while serving `ranges`. +/// +/// The merge rule bounds one read, not a whole fetch: every merged range comes back as its own +/// buffer and all of them are live together. With the aligned cache, the aligned copies and any +/// range assembled from several of them coexist as well. +fn fetch_peak( + reader: &PayloadReader<'_>, + ranges: &[std::ops::Range], + area_end: u64, + aligned: bool, +) -> usize { + if !aligned { + return merge_read_ranges(ranges) + .iter() + .map(|range| (range.end - range.start) as usize) + .sum(); + } + reader.reservation_bytes(ranges, area_end) + + ranges + .iter() + .map(|range| (range.end - range.start) as usize) + .sum::() +} + +/// Fetches the exact keys of the blocks of one metadata page that survive coarse selection. +/// +/// A descriptor only references its key, so identity never costs a copy per row or per block +/// descriptor. Keys are read only for selected blocks and at most once per series, and the +/// selected blocks of a page are merged into runs bounded by [`MAX_READ_BYTES`], so a page of +/// many short series costs one read rather than one per series. +/// +/// The pool is charged the coexisting peak before the read - the fetched or aligned and assembled +/// bytes plus the owned keys decoded from them - and shrunk back to the retained keys afterwards. +/// A pool that cannot pay the aligned peak is asked for the unbuffered one, then for a narrower +/// window; the keys left out are read one at a time, which is where a key that does not fit at all +/// reports exhaustion. Reading first and refusing the key afterwards would leave the bytes already +/// fetched. +#[allow(clippy::too_many_arguments)] +async fn load_page_keys( + reader: &PayloadReader<'_>, + areas: Areas, + blocks: &[Block], + request: &SamplesReadRequest, + ids: Option<&BTreeSet>, + keys: &mut HashMap<(u64, u32), ExactSeriesKey>, + reservation: Option<&MemoryReservation>, + offsets: &mut mito_codec::row_converter::SparseOffsetsCache, + metrics: &mut ColumnarReadMetrics, +) -> Result<()> { + // Selected blocks whose key is not cached yet, one per reference: blocks of one series + // share a copy, and blocks are ordered by series, so the references increase. + let mut wanted: Vec<&Block> = Vec::new(); + for block in blocks { + if !matches(block.bounds(), request, ids) + || keys.contains_key(&(block.key.offset, block.key.length)) + { + continue; + } + block.key.validate(areas.keys)?; + if wanted.last().is_some_and(|last| last.key == block.key) { + continue; + } + wanted.push(block); + } + if wanted.is_empty() { + return Ok(()); + } + // The same range fetch the column chunks use: aligned blocks of the part are shared with the + // payload cache, so a part that is already cached costs no key read at all. Alignment is the + // first choice and the first victim: it costs a whole aligned block per key area touched. + let uncached = PayloadReader { + reader: reader.reader, + cache: None, + end: reader.end, + key_end: reader.key_end, + }; + let mut window = wanted.len(); + let mut attempt = None; + while window > 0 { + let selected = &wanted[..window]; + let ranges = selected + .iter() + .map(|block| Ok(block.key.offset..block.key.checksum_end()?)) + .collect::>>()?; + let retained = selected + .iter() + .map(|block| retained_key_bytes(block.key.length)) + .sum::(); + let mut candidates = Vec::with_capacity(2); + if reader.cache.is_some() { + candidates.push(true); + } + candidates.push(false); + for aligned in candidates { + // The fetched entries and the owned keys they become are live together; the fetch + // peak already covers the entries themselves. + let peak = fetch_peak(reader, &ranges, areas.keys.end, aligned) + retained; + if try_reserve_keys(reservation, peak) { + attempt = Some((aligned, peak, retained, ranges)); + break; + } + } + if attempt.is_some() { + break; + } + // Halve the window rather than give up: the keys left out keep their own, smaller, + // on-demand reservation. + window /= 2; + } + let Some((aligned, peak, retained, ranges)) = attempt else { + return Ok(()); + }; + let _peak = KeyPeak { + reservation, + transient: peak - retained, + }; + let target = if aligned { reader } else { &uncached }; + let fetched = fetch_ranges(target, areas.keys.end, &ranges, metrics).await?; + for (block, entry) in wanted.iter().take(window).zip(fetched) { + let key = + ExactSeriesKey::from_sparse_key(verified_key(&entry, block.key.length)?, offsets)?; + // A key that parses but claims another identifier is a malformed reference: ownership + // would silently point at the wrong series. + require( + key.id() == block.id, + "sample block series key disagrees with its identifier", + )?; + keys.insert((block.key.offset, block.key.length), key); + } + Ok(()) +} + +/// Bounds-only pruning. Selection itself is exact and decided by [`owns`]. +fn matches( + bounds: Bounds, + request: &SamplesReadRequest, + ids: Option<&BTreeSet>, +) -> bool { if request.start.is_some_and(|start| bounds.max < start) || request.end.is_some_and(|end| bounds.min >= end) { @@ -328,11 +802,7 @@ fn matches(bounds: Bounds, request: &SamplesReadRequest) -> bool { tsid: u64::MAX, }); } - first <= last - && request - .ids - .as_ref() - .is_none_or(|ids| ids.range(first..=last).next().is_some()) + first <= last && ids.is_none_or(|ids| ids.range(first..=last).next().is_some()) } async fn read_range(store: &ObjectStore, path: &str, start: u64, end: u64) -> Result { require(start <= end, "reversed sample byte range")?; @@ -355,13 +825,13 @@ async fn load_metadata_page( object: (&str, u64), reference: &PageReference, depth: u32, - payload_end: u64, + areas: Areas, manifest: &Manifest, ) -> Result> { let (path, size) = object; metadata_cache::load( cache, - format!("{path}/{size}/{reference:?}/{depth}/{payload_end}"), + format!("{path}/{size}/{reference:?}/{depth}/{}", areas.payload_end), Page::retained_size, async { if reference.payload.encoding != 1 { @@ -373,7 +843,7 @@ async fn load_metadata_page( &reference.payload.unpack(&bytes)?, reference, depth, - payload_end, + areas, manifest, ) }, @@ -401,6 +871,16 @@ fn load_column_chunk<'a>( .ok_or_else(|| invalid("missing prefetched sample column"))?; column.payload.unpack_with_decoder(bytes, decoder) } +/// Fixed per-column charge `arrow` reports beyond the column's own data. +/// +/// `Array::get_array_memory_size` is `size_of::()` plus the buffer capacities, and +/// `RecordBatch::get_array_memory_size` sums those with no other term. `BinaryArray` is the +/// widest array type this reader can build, so its struct size bounds every projected column. +const ARRAY_HEADER_BYTES: usize = std::mem::size_of::(); +/// Granularity `arrow_buffer` rounds every buffer allocation up to, so each column can retain up +/// to this much more than the bytes it holds. +const BUFFER_ALIGNMENT: usize = 64; + fn timestamp_array(values: impl Iterator, unit: TimestampUnit) -> ArrayRef { match unit { TimestampUnit::Second => Arc::new(TimestampSecondArray::from_iter_values(values)), @@ -450,14 +930,25 @@ async fn read_block( .map(|r| r.end - r.start) .sum::() - ranges.iter().map(|r| r.end - r.start).sum::(); + // The block's exact key is returned with the batch, once per block rather than once per + // row, so it costs one key regardless of how many rows the block holds. Time pruning + // cannot add output rows, so no per-row term is needed for identity at all. + let key_bytes = + std::mem::size_of::() + block.key.length as usize + BUFFER_ALIGNMENT; let base_bytes = 2 * gap_bytes as usize + columns .iter() .map(|column| column.payload.stored as usize + column.payload.decoded as usize) .sum::() + block.rows as usize * (32 * (projection.len() + 1) + std::mem::size_of::()) + + key_bytes + // Each projected column also retains one array header and one allocation rounded up to + // the allocator's granularity, which the per-row bytes above do not cover. A block with + // one row and an empty projection selects no payload column at all, so this is the only + // term that keeps its array accounted for. + + projection.len() * (ARRAY_HEADER_BYTES + BUFFER_ALIGNMENT) + columns.len() * 128; - let bytes = base_bytes + reader.reservation_bytes(&ranges); + let bytes = base_bytes + reader.reservation_bytes(&ranges, reader.end); let mut bypass_cache = false; if let Some(reservation) = &reservation { let mut result = reservation.try_grow(bytes); @@ -483,6 +974,7 @@ async fn read_block( reader: reader.reader, cache: None, end: reader.end, + key_end: reader.key_end, }; let reader = if bypass_cache { &uncached } else { reader }; let mut buffers = HashMap::new(); @@ -540,61 +1032,71 @@ async fn read_block( // Visit columns in physical order, regardless of requested output order. let mut sorted = projection.to_vec(); sorted.sort_unstable(); + // The exact key column is never projected: it is returned with the batch instead, + // so a block of any row count carries one key in total. for index in sorted { - let array: ArrayRef = match index { - 0 => Arc::new(UInt32Array::from(vec![block.id.table_id; selected.len()])), - 1 => Arc::new(UInt64Array::from(vec![block.id.tsid; selected.len()])), - 2 => { - let values = times - .as_ref() - .ok_or_else(|| invalid("missing decoded sample time"))?; - timestamp_array( - selected.iter().map(|&i| values[i] as i64), - manifest.timestamp.unit, - ) - } - 3 => { - let column = &block.columns[1]; - let values = codec::decode_numbers( - load_column_chunk(&buffers, column, &mut decoder)?, - column.payload.encoding, - block.rows as usize, - )?; - if let Some(times) = × { - require( - (1..values.len()) - .all(|i| times[i - 1] != times[i] || values[i - 1] >= values[i]), - "unordered sample sequences", - )?; + let array: ArrayRef = if index == 0 { + Arc::new(UInt32Array::from(vec![block.id.table_id; selected.len()])) + } else if index == 1 { + Arc::new(UInt64Array::from(vec![block.id.tsid; selected.len()])) + } else { + match index { + 2 => { + let values = times + .as_ref() + .ok_or_else(|| invalid("missing decoded sample time"))?; + timestamp_array( + selected.iter().map(|&i| values[i] as i64), + manifest.timestamp.unit, + ) } - Arc::new(UInt64Array::from_iter_values( - selected.iter().map(|&i| values[i]), - )) - } - 4 => { - let column = &block.columns[2]; - let values = codec::decode_operations( - load_column_chunk(&buffers, column, &mut decoder)?, - column.payload.encoding, - block.rows as usize, - )?; - Arc::new(UInt8Array::from_iter_values( - selected.iter().map(|&i| values[i]), - )) - } - field => { - let id = manifest.fields[field - 5].column_id; - if let Some(column) = block.columns.iter().find(|c| c.role == 4 && c.id == id) { - let values = codec::decode_float( + 3 => { + let column = &block.columns[1]; + let values = codec::decode_numbers( load_column_chunk(&buffers, column, &mut decoder)?, + column.payload.encoding, block.rows as usize, - column.nulls, )?; - Arc::new(Float64Array::from_iter( - selected.iter().map(|&i| values[i].map(f64::from_bits)), + if let Some(times) = × { + require( + (1..values.len()) + .all(|i| times[i - 1] != times[i] || values[i - 1] >= values[i]), + "unordered sample sequences", + )?; + } + Arc::new(UInt64Array::from_iter_values( + selected.iter().map(|&i| values[i]), )) - } else { - Arc::new(Float64Array::from(vec![None::; selected.len()])) + } + 4 => { + let column = &block.columns[2]; + let values = codec::decode_operations( + load_column_chunk(&buffers, column, &mut decoder)?, + column.payload.encoding, + block.rows as usize, + )?; + Arc::new(UInt8Array::from_iter_values( + selected.iter().map(|&i| values[i]), + )) + } + field => { + let field = manifest + .fields + .get(field - 5) + .ok_or_else(|| invalid("unknown sample field"))?; + let id = field.column_id; + if let Some(column) = block.columns.iter().find(|c| c.role == 4 && c.id == id) { + let values = codec::decode_float( + load_column_chunk(&buffers, column, &mut decoder)?, + block.rows as usize, + column.nulls, + )?; + Arc::new(Float64Array::from_iter( + selected.iter().map(|&i| values[i].map(f64::from_bits)), + )) + } else { + Arc::new(Float64Array::from(vec![None::; selected.len()])) + } } } }; @@ -653,106 +1155,76 @@ async fn fetch_columns( .iter() .map(|column| Ok(column.payload.offset..column.payload.end()?)) .collect::>>()?; - let merged = merge_read_ranges(&ranges); - *metrics.counts.entry("column_ranges").or_default() += ranges.len() as u64; - *metrics.bytes.entry("requested").or_default() += - ranges.iter().map(|r| r.end - r.start).sum::(); - if reader.cache.is_some() { - return fetch_cached_columns(reader, &columns, &ranges, buffers, metrics).await; + let fetched = fetch_ranges(reader, reader.end, &ranges, metrics).await?; + for (column, bytes) in columns.iter().zip(fetched) { + buffers.insert(column.payload.offset, bytes); } - *metrics.counts.entry("payload_reads").or_default() += merged.len() as u64; - let started = Instant::now(); - let fetched = reader - .reader - .fetch(merged.clone()) - .await - .context(OpenDalSnafu)?; - metrics.add_elapsed("columnar_payload_fetch_wall", started.elapsed()); - require( - fetched.len() == merged.len(), - "sample fetch response count mismatch", - )?; - let mut column_index = 0; - for (range, buffer) in merged.iter().zip(fetched) { - let bytes = buffer.to_bytes(); - require( - bytes.len() as u64 == range.end - range.start, - "sample range length mismatch", - )?; - *metrics.bytes.entry("fetched").or_default() += bytes.len() as u64; - while let Some(column) = columns.get(column_index) { - if column.payload.offset >= range.end { - break; - } - require( - column.payload.offset >= range.start && column.payload.end()? <= range.end, - "sample column outside fetched range", - )?; - let start = (column.payload.offset - range.start) as usize; - let end = (column.payload.end()? - range.start) as usize; - require(end <= bytes.len(), "sample column outside fetched range")?; - buffers.insert(column.payload.offset, bytes.slice(start..end)); - column_index += 1; - } - } - require( - column_index == columns.len(), - "missing sample column response", - )?; Ok(()) } -/// Stable cache units; adjacent missing units still use bounded merged reads. -const PAYLOAD_BLOCK_BYTES: u64 = 16 * 1024; - -struct PayloadReader<'a> { - reader: &'a Reader, - cache: Option, - end: u64, -} - -impl PayloadReader<'_> { - fn reservation_bytes(&self, ranges: &[std::ops::Range]) -> usize { - if self.cache.is_none() { - return 0; - } - // Aligned read allocations and assembly can coexist. Cache-owned copies - // are separately charged to the shared page-cache budget. - 2 * aligned_payload_ranges(ranges, self.end) - .iter() - .map(|r| (r.end - r.start) as usize) - .sum::() - } -} - -fn aligned_payload_ranges(ranges: &[std::ops::Range], end: u64) -> Vec> { - let mut offsets = std::collections::BTreeSet::new(); - for range in ranges { - if range.start >= range.end { - continue; - } - for block in range.start / PAYLOAD_BLOCK_BYTES..=(range.end - 1) / PAYLOAD_BLOCK_BYTES { - offsets.insert(block * PAYLOAD_BLOCK_BYTES); - } - } - offsets - .into_iter() - .map(|start| start..start.saturating_add(PAYLOAD_BLOCK_BYTES).min(end)) - .collect() -} - -async fn fetch_cached_columns( +/// Fetch arbitrary byte ranges of one part, reusing its aligned block cache when it has one. +/// +/// `end` is where the ranges' area ends (the payload area or the key area), so an aligned block +/// never extends past the area it belongs to. Keys use the same path as column chunks: a part +/// whose blocks are already cached issues no read at all, which is why a second series in the +/// same part can be served from the over-read that the first one paid for. +async fn fetch_ranges( reader: &PayloadReader<'_>, - columns: &[&Column], + end: u64, ranges: &[std::ops::Range], - buffers: &mut HashMap, metrics: &mut ColumnarReadMetrics, -) -> Result<()> { +) -> Result> { + *metrics.counts.entry("column_ranges").or_default() += ranges.len() as u64; + *metrics.bytes.entry("requested").or_default() += + ranges.iter().map(|r| r.end - r.start).sum::(); + if ranges.is_empty() { + return Ok(Vec::new()); + } let Some(cache) = &reader.cache else { - return Err(invalid("missing payload cache")); + // Merging is bounded by [`MAX_READ_BYTES`], and a requested range never exceeds that, so + // every range lies inside exactly one merged read. + let merged = merge_read_ranges(ranges); + *metrics.counts.entry("payload_reads").or_default() += merged.len() as u64; + let started = Instant::now(); + let fetched = reader + .reader + .fetch(merged.clone()) + .await + .context(OpenDalSnafu)?; + metrics.add_elapsed("columnar_payload_fetch_wall", started.elapsed()); + require( + fetched.len() == merged.len(), + "sample fetch response count mismatch", + )?; + let mut result = Vec::with_capacity(ranges.len()); + let mut index = 0; + for (range, buffer) in merged.iter().zip(fetched) { + let bytes = buffer.to_bytes(); + require( + bytes.len() as u64 == range.end - range.start, + "sample range length mismatch", + )?; + *metrics.bytes.entry("fetched").or_default() += bytes.len() as u64; + while let Some(requested) = ranges.get(index) { + if requested.start >= range.end { + break; + } + require( + requested.start >= range.start && requested.end <= range.end, + "sample range outside fetched range", + )?; + result.push(bytes.slice( + (requested.start - range.start) as usize + ..(requested.end - range.start) as usize, + )); + index += 1; + } + } + require(index == ranges.len(), "missing sample range response")?; + return Ok(result); }; - let aligned = aligned_payload_ranges(ranges, reader.end); - let mut blocks = HashMap::new(); + let aligned = aligned_ranges(ranges, end); + let mut blocks: HashMap = HashMap::new(); let mut missing = Vec::new(); for range in &aligned { if let Some(bytes) = cache.get(range) { @@ -766,8 +1238,6 @@ async fn fetch_cached_columns( if !merged.is_empty() { *metrics.counts.entry("payload_reads").or_default() += merged.len() as u64; let started = Instant::now(); - // No shared guard or pending-owner dependency crosses this await. A - // backpressured or cancelled stream cannot stall another partition. let fetched = reader .reader .fetch(merged.clone()) @@ -804,7 +1274,8 @@ async fn fetch_cached_columns( } require(missing.next().is_none(), "missing aligned block response")?; } - for (column, range) in columns.iter().zip(ranges) { + let mut result = Vec::with_capacity(ranges.len()); + for range in ranges { let start = range.start / PAYLOAD_BLOCK_BYTES * PAYLOAD_BLOCK_BYTES; let first = blocks .get(&start) @@ -829,9 +1300,53 @@ async fn fetch_cached_columns( } assembled.freeze() }; - buffers.insert(column.payload.offset, bytes); + result.push(bytes); } - Ok(()) + Ok(result) +} + +/// Stable cache units; adjacent missing units still use bounded merged reads. +const PAYLOAD_BLOCK_BYTES: u64 = 16 * 1024; + +struct PayloadReader<'a> { + reader: &'a Reader, + cache: Option, + /// End of the column payload area; aligned payload blocks are clipped here. + end: u64, + /// End of the key area; aligned key blocks are clipped here. + key_end: u64, +} + +impl PayloadReader<'_> { + /// Aligned and assembled bytes that serving `ranges` inside `area_end` can hold at once. + fn reservation_bytes(&self, ranges: &[std::ops::Range], area_end: u64) -> usize { + if self.cache.is_none() { + return 0; + } + // Aligned read allocations and assembly can coexist. Cache-owned copies + // are separately charged to the shared page-cache budget. + 2 * aligned_ranges(ranges, area_end) + .iter() + .map(|r| (r.end - r.start) as usize) + .sum::() + } +} + +/// Aligned cache units covering `ranges`, clipped at `end`. +fn aligned_ranges(ranges: &[std::ops::Range], end: u64) -> Vec> { + let mut offsets = std::collections::BTreeSet::new(); + for range in ranges { + if range.start >= range.end { + continue; + } + for block in range.start / PAYLOAD_BLOCK_BYTES..=(range.end - 1) / PAYLOAD_BLOCK_BYTES { + offsets.insert(block * PAYLOAD_BLOCK_BYTES); + } + } + offsets + .into_iter() + .map(|start| start..start.saturating_add(PAYLOAD_BLOCK_BYTES).min(end)) + .collect() } const MAX_READ_GAP: u64 = 4 * 1024; @@ -878,6 +1393,7 @@ async fn prefetch_blocks( start: usize, field_ids: &HashSet, request: &SamplesReadRequest, + ids: Option<&BTreeSet>, projection: &[usize], payload_end: u64, metrics: &mut ColumnarReadMetrics, @@ -889,7 +1405,7 @@ async fn prefetch_blocks( let mut payload_ranges = Vec::new(); let mut end = start; for block in blocks.iter().skip(start).take(MAX_BLOCKS) { - if !matches(block.bounds(), request) { + if !matches(block.bounds(), request, ids) { end += 1; continue; } @@ -935,7 +1451,7 @@ async fn prefetch_blocks( payload_ranges.push(column.payload.offset..column.payload.end()?); } let fetch_bytes = if reader.cache.is_some() { - aligned_payload_ranges(&payload_ranges, reader.end) + aligned_ranges(&payload_ranges, reader.end) .iter() .map(|range| (range.end - range.start) as usize) .sum() @@ -968,6 +1484,7 @@ async fn prefetch_blocks( .iter() .map(|c| Ok(c.payload.offset..c.payload.end()?)) .collect::>>()?, + reader.end, ), ) .is_err() diff --git a/src/mito2/src/series_index/samples/tests.rs b/src/mito2/src/series_index/samples/tests.rs index 92900647872..6a5bf0b0257 100644 --- a/src/mito2/src/series_index/samples/tests.rs +++ b/src/mito2/src/series_index/samples/tests.rs @@ -31,10 +31,13 @@ use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::FileId; -use super::metadata::{Bounds, Page, Trailer}; +use super::metadata::{ + Bounds, HEADER_SIZE, KEY_CHECKSUM_BYTES, Page, PayloadReference, TRAILER_SIZE, Trailer, +}; use super::*; use crate::series_index::artifact::{SamplesEntry, SamplesFileHandle}; -use crate::series_index::format::{SamplesFormat, SamplesReadRequest}; +use crate::series_index::format::{SamplesBatch, SamplesFormat, SamplesReadRequest}; +use crate::series_index::identity::{ExactSeriesKey, IdentityEncoding}; use crate::series_index::purger::series_index_channel; use crate::sst::file::RegionFileId; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; @@ -58,9 +61,48 @@ fn batch(meta: &RegionMetadataRef, rows: &[Sample], dictionary: bool) -> RecordB .iter() .map(|r| new_sparse_primary_key(&["a", "b"], meta, r.0, r.1)) .collect(); + sample_batch( + meta, + &keys, + rows.iter().map(|r| r.2).collect(), + rows.iter().map(|r| r.3).collect(), + rows.iter().map(|r| r.4).collect(), + rows.iter().map(|r| r.5.map(f64::from_bits)).collect(), + dictionary, + ) +} + +/// Builds one sample batch from explicit sparse keys, so a test can give two +/// distinct series the same `tsid` and different labels. +fn keyed_batch( + meta: &RegionMetadataRef, + keys: &[Vec], + times: &[i64], + dictionary: bool, +) -> RecordBatch { + sample_batch( + meta, + keys, + times.to_vec(), + vec![1; keys.len()], + vec![0; keys.len()], + vec![Some(0.0); keys.len()], + dictionary, + ) +} + +fn sample_batch( + meta: &RegionMetadataRef, + keys: &[Vec], + times: Vec, + sequences: Vec, + operations: Vec, + values: Vec>, + dictionary: bool, +) -> RecordBatch { let keys: ArrayRef = if dictionary { let mut builder = BinaryDictionaryBuilder::::new(); - for key in &keys { + for key in keys { builder.append(key).unwrap(); } Arc::new(builder.finish()) @@ -77,26 +119,18 @@ fn batch(meta: &RegionMetadataRef, rows: &[Sample], dictionary: bool) -> RecordB .unwrap() .unit(); let ts: ArrayRef = match unit { - common_time::timestamp::TimeUnit::Second => Arc::new( - datatypes::arrow::array::TimestampSecondArray::from_iter_values( - rows.iter().map(|r| r.2), - ), - ), - common_time::timestamp::TimeUnit::Millisecond => Arc::new( - datatypes::arrow::array::TimestampMillisecondArray::from_iter_values( - rows.iter().map(|r| r.2), - ), - ), - common_time::timestamp::TimeUnit::Microsecond => Arc::new( - datatypes::arrow::array::TimestampMicrosecondArray::from_iter_values( - rows.iter().map(|r| r.2), - ), - ), - common_time::timestamp::TimeUnit::Nanosecond => Arc::new( - datatypes::arrow::array::TimestampNanosecondArray::from_iter_values( - rows.iter().map(|r| r.2), - ), - ), + common_time::timestamp::TimeUnit::Second => { + Arc::new(datatypes::arrow::array::TimestampSecondArray::from_iter_values(times)) + } + common_time::timestamp::TimeUnit::Millisecond => { + Arc::new(datatypes::arrow::array::TimestampMillisecondArray::from_iter_values(times)) + } + common_time::timestamp::TimeUnit::Microsecond => { + Arc::new(datatypes::arrow::array::TimestampMicrosecondArray::from_iter_values(times)) + } + common_time::timestamp::TimeUnit::Nanosecond => { + Arc::new(datatypes::arrow::array::TimestampNanosecondArray::from_iter_values(times)) + } }; RecordBatch::try_new( Arc::new(Schema::new(vec![ @@ -107,13 +141,11 @@ fn batch(meta: &RegionMetadataRef, rows: &[Sample], dictionary: bool) -> RecordB Field::new(OPERATION, DataType::UInt8, false), ])), vec![ - Arc::new(Float64Array::from_iter( - rows.iter().map(|r| r.5.map(f64::from_bits)), - )), + Arc::new(Float64Array::from_iter(values)), ts, keys, - Arc::new(UInt64Array::from_iter_values(rows.iter().map(|r| r.3))), - Arc::new(UInt8Array::from_iter_values(rows.iter().map(|r| r.4))), + Arc::new(UInt64Array::from_iter_values(sequences)), + Arc::new(UInt8Array::from_iter_values(operations)), ], ) .unwrap() @@ -148,6 +180,40 @@ async fn build( purger, ) } +/// Reads the descriptor and root page of a one-part artifact, plus the part's bytes. Tests +/// use it to inspect what the reader decodes and to rewrite a stored page. +async fn root_page( + store: &ObjectStore, + handle: &SamplesFileHandle, +) -> (Manifest, Trailer, Page, Vec) { + let manifest = manifest(store, handle).await; + assert_eq!( + 1, + manifest.files.len(), + "one part keeps the root page addressable" + ); + let bytes = store + .read(&manifest.part_path(0)) + .await + .unwrap() + .to_bytes() + .to_vec(); + let trailer_start = bytes.len() - TRAILER_SIZE as usize; + let trailer = Trailer::decode(&bytes[trailer_start..], bytes.len() as u64).unwrap(); + assert_eq!(1, trailer.depth, "one leaf keeps the root page addressable"); + let start = trailer.root.payload.offset as usize; + let stored = &bytes[start..start + trailer.root.payload.stored as usize]; + let page = Page::decode( + &trailer.root.payload.unpack(stored).unwrap(), + &trailer.root, + trailer.depth, + trailer.areas(), + &manifest, + ) + .unwrap(); + (manifest, trailer, page, bytes) +} + async fn manifest(store: &ObjectStore, handle: &SamplesFileHandle) -> Manifest { let path = &handle .entry() @@ -175,7 +241,7 @@ fn full_request() -> SamplesReadRequest { "field_0", ]) } -fn collect_rows(batches: &[RecordBatch]) -> Vec { +fn collect_rows(batches: &[SamplesBatch]) -> Vec { batches .iter() .flat_map(|b| { @@ -336,7 +402,8 @@ async fn equal_timestamps_cross_leaves_and_parts() { let m = metadata(); let format = ColumnarSamplesFormat { block_rows: 1, - page_bytes: 232, + // One block descriptor plus header; a block now also stores its exact series key. + page_bytes: 288, root_capacity: 2, ..Default::default() }; @@ -349,18 +416,18 @@ async fn equal_timestamps_cross_leaves_and_parts() { assert!(manifest.files.len() >= 4); let first = store.read(&manifest.part_path(0)).await.unwrap().to_bytes(); assert_eq!( - Trailer::decode(&first[first.len() - 112..], first.len() as u64) - .unwrap() - .depth, + Trailer::decode( + &first[first.len() - TRAILER_SIZE as usize..], + first.len() as u64 + ) + .unwrap() + .depth, 2 ); let reader = format.reader(store, handle).await.unwrap(); for (start, end) in [(99, 100), (100, 101), (101, 102), (50, 101)] { let mut req = full_request(); - req.ids = Some(BTreeSet::from([MetricSeriesId { - table_id: 7, - tsid: 42, - }])); + req.keys = Some(BTreeSet::from([exact(&m, &["a", "b"], 7, 42)])); req.start = Some(start); req.end = Some(end); let out = reader @@ -388,6 +455,288 @@ async fn equal_timestamps_cross_leaves_and_parts() { assert_eq!(out.len(), 1); assert_eq!(out[0].column(0).null_count(), 1); } +/// Wraps one encoded primary key in the canonical exact identity. +fn exact(meta: &RegionMetadataRef, tags: &[&str; 2], table_id: u32, tsid: u64) -> ExactSeriesKey { + let mut offsets = mito_codec::row_converter::SparseOffsetsCache::new(); + ExactSeriesKey::from_sparse_key( + &new_sparse_primary_key(tags, meta, table_id, tsid), + &mut offsets, + ) + .unwrap() +} + +/// Reads the samples selected by one exact series key. +async fn read_key(reader: &dyn SamplesFormatReader, key: ExactSeriesKey) -> Vec { + let mut req = full_request(); + req.keys = Some(BTreeSet::from([key])); + let out = reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap(); + collect_rows(&out) +} + +/// A key whose label is long enough to push the encoded key past 8 KiB round-trips like any +/// other: the writer and the reader agree because neither imposes a length cap. +#[tokio::test] +async fn long_keys_round_trip() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(8 * 1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 7); + assert!(key.len() > 8 * 1024, "{}", key.len()); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[key.clone(), key.clone()], + &[100, 200], + false, + )], + ) + .await; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let selection = exact(&meta, &[label.as_str(), "z"], 1, 7); + assert_eq!(selection.bytes(), key.as_slice()); + assert_eq!( + read_key(reader.as_ref(), selection).await, + vec![(1, 7, 100, 1, 0, Some(0)), (1, 7, 200, 1, 0, Some(0)),] + ); + // The key area keeps the key byte for byte, so a read returns the same bytes the writer + // received; the descriptor only references it. + let (_, trailer, page, bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!(1, blocks.len()); + let reference = blocks[0].key; + reference.validate(trailer.keys()).unwrap(); + let end = reference.offset + u64::from(reference.length); + assert_eq!( + key.as_slice(), + &bytes[reference.offset as usize..end as usize] + ); +} + +/// The captured pair in the samples artifact: two distinct series with one +/// `table_id`/`tsid` whose timestamps restart. Every block records which exact key +/// owns it, so each series keeps its own samples and never sees the partner's. +#[tokio::test] +async fn colliding_tsids_are_retained_and_owned_exactly() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["host-020078", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-029838", "z"], &meta, 1, 7); + for dictionary in [false, true] { + let handle = build( + &format, + &store, + meta.clone(), + &[ + keyed_batch( + &meta, + &[first.clone(), first.clone()], + &[100, 200], + dictionary, + ), + keyed_batch( + &meta, + &[partner.clone(), partner.clone()], + &[100, 200], + dictionary, + ), + ], + ) + .await; + let reader = format.reader(store.clone(), handle).await.unwrap(); + let first_key = exact(&meta, &["host-020078", "z"], 1, 7); + let partner_key = exact(&meta, &["host-029838", "z"], 1, 7); + assert_eq!( + read_key(reader.as_ref(), first_key.clone()).await, + vec![(1, 7, 100, 1, 0, Some(0)), (1, 7, 200, 1, 0, Some(0)),] + ); + assert_eq!( + read_key(reader.as_ref(), partner_key.clone()).await, + vec![(1, 7, 100, 1, 0, Some(0)), (1, 7, 200, 1, 0, Some(0)),] + ); + // Selecting both owners reads both, even though they share one identifier. + let mut req = full_request(); + req.keys = Some(BTreeSet::from([first_key.clone(), partner_key.clone()])); + let out = reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_rows(&out).len(), 4); + // Selecting one owner never reads the partner's block. + let mut req = full_request(); + req.keys = Some(BTreeSet::from([partner_key])); + let out = reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_rows(&out).len(), 2); + } +} + +/// The silent-loss variant: timestamps keep increasing across the collision, so the +/// old writer accepted it and both series' samples landed under one identity. +#[tokio::test] +async fn colliding_tsids_with_increasing_timestamps_are_retained() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["host-020078", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-029838", "z"], &meta, 1, 7); + for dictionary in [false, true] { + let handle = build( + &format, + &store, + meta.clone(), + &[ + keyed_batch( + &meta, + &[first.clone(), first.clone()], + &[100, 200], + dictionary, + ), + keyed_batch( + &meta, + &[partner.clone(), partner.clone()], + &[300, 400], + dictionary, + ), + ], + ) + .await; + let reader = format.reader(store.clone(), handle).await.unwrap(); + let first_key = exact(&meta, &["host-020078", "z"], 1, 7); + let partner_key = exact(&meta, &["host-029838", "z"], 1, 7); + assert_eq!( + read_key(reader.as_ref(), first_key).await, + vec![(1, 7, 100, 1, 0, Some(0)), (1, 7, 200, 1, 0, Some(0)),] + ); + assert_eq!( + read_key(reader.as_ref(), partner_key).await, + vec![(1, 7, 300, 1, 0, Some(0)), (1, 7, 400, 1, 0, Some(0)),] + ); + } +} + +/// Many distinct keys under one forced tsid, through both primary-key inputs. +#[tokio::test] +async fn many_distinct_keys_under_one_tsid_are_retained() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let keys: Vec<_> = (0..8) + .map(|index| { + let label = format!("host-{index:06}"); + new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 5) + }) + .collect(); + let times: Vec<_> = (0..8).map(|index| 100 + index as i64).collect(); + for dictionary in [false, true] { + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &keys, ×, dictionary)], + ) + .await; + let reader = format.reader(store.clone(), handle).await.unwrap(); + let all = reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_rows(&all).len(), keys.len()); + // Each owner sees exactly one row, at its own timestamp. + for index in 0..8 { + let label = format!("host-{index:06}"); + let key = exact(&meta, &[label.as_str(), "z"], 1, 5); + let rows = read_key(reader.as_ref(), key).await; + assert_eq!(rows, vec![(1, 5, 100 + index as i64, 1, 0, Some(0))]); + } + // An absent owner selects nothing, even though its identifier is present. + let absent = exact(&meta, &["host-999999", "z"], 1, 5); + assert!(read_key(reader.as_ref(), absent).await.is_empty()); + } +} + +/// A regression in sparse-key order is refused; the same rules that guard one series +/// still apply to a stream of distinct series. +#[tokio::test] +async fn out_of_order_series_are_refused() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["a", "b"], &meta, 1, 1); + let second = new_sparse_primary_key(&["a", "b"], &meta, 1, 2); + let file_id = FileId::random(); + let generation = FileId::random(); + let prefix = format!("{}/samples/{generation}/", meta.region_id.as_u64()); + let mut writer = format + .writer(meta.clone(), store.clone(), &prefix, file_id, generation) + .await + .unwrap(); + writer + .write(&keyed_batch( + &meta, + &[first.clone(), second.clone()], + &[100, 200], + false, + )) + .await + .unwrap(); + let error = writer + .write(&keyed_batch(&meta, &[first], &[300], false)) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("not sorted by primary key"), "{error}"); + writer.abort().await.unwrap(); + // Timestamps are still monotonic inside one exact series. + let generation = FileId::random(); + let prefix = format!("{}/samples/{generation}/", meta.region_id.as_u64()); + let mut writer = format + .writer( + meta.clone(), + store.clone(), + &prefix, + FileId::random(), + generation, + ) + .await + .unwrap(); + writer + .write(&keyed_batch( + &meta, + std::slice::from_ref(&second), + &[200], + false, + )) + .await + .unwrap(); + let error = writer + .write(&keyed_batch(&meta, &[second], &[100], false)) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("sample input is not sorted"), "{error}"); + writer.abort().await.unwrap(); +} + #[tokio::test] async fn empty_selection_and_validation() { let store = store(); @@ -413,7 +762,7 @@ async fn empty_selection_and_validation() { req.end = Some(1); assert!(reader.read(req).is_err()); let mut req = request(&[]); - req.ids = Some(BTreeSet::new()); + req.keys = Some(BTreeSet::new()); assert!( reader .read(req) @@ -551,8 +900,9 @@ async fn selective_reads_pin_the_artifact() { assert_eq!(result.len(), 1); assert_eq!(result[0].column(0).null_count(), 1); assert!(purges.try_recv().is_ok()); - // Header, trailer, one leaf/root. The all-null projected field has no payload. - assert_eq!(reads.lock().unwrap().len(), 3); + // Header, trailer, one leaf/root, and the selected block's key from the key area. The + // all-null projected field has no payload. + assert_eq!(reads.lock().unwrap().len(), 4); } #[test] fn payload_limits_and_timestamp_envelopes() { @@ -611,7 +961,7 @@ async fn malformed_framing_and_payload_corruption() { let manifest = manifest(&store, &handle).await; let path = manifest.part_path(0); let original = store.read(&path).await.unwrap().to_vec(); - let end = original.len() - 112; + let end = original.len() - TRAILER_SIZE as usize; let trailer = Trailer::decode(&original[end..], original.len() as u64).unwrap(); for (encoding, compression) in [(99, 0), (1, 99)] { let mut unsupported = Trailer::decode(&original[end..], original.len() as u64).unwrap(); @@ -625,7 +975,7 @@ async fn malformed_framing_and_payload_corruption() { .unpack(&original[trailer.root.payload.offset as usize..end]) .unwrap(); let Page::Leaf(blocks) = - Page::decode(&decoded, &trailer.root, 1, trailer.payload_end, &manifest).unwrap() + Page::decode(&decoded, &trailer.root, 1, trailer.areas(), &manifest).unwrap() else { panic!() }; @@ -649,7 +999,7 @@ async fn malformed_framing_and_payload_corruption() { .await .is_err() ); - for offset in [0, end + 72, end + 104] { + for offset in [0, end + 72, end + TRAILER_SIZE as usize - 8] { let mut bad = original.clone(); bad[offset] ^= 3; store.write(&path, bad).await.unwrap(); @@ -668,7 +1018,7 @@ async fn malformed_framing_and_payload_corruption() { &decoded[..length], &trailer.root, 1, - trailer.payload_end, + trailer.areas(), &manifest ) .is_err() @@ -845,7 +1195,151 @@ async fn real_samples_reconcile_reopen_and_retire( engine.stop().await.unwrap(); } -async fn fixture(store: &ObjectStore) -> SamplesFileHandle { +/// An artifact written by the previous layout keeps its key in every block descriptor and names +/// that older encoding in its manifest. +/// +/// Nothing else separates it from a current artifact: it is structurally valid, ordered, and has +/// the same version. Without the layout in `identity_encoding`, maintenance opens it, finds it +/// sound, and keeps it - and only the first part read fails. The manifest must be refused instead, +/// so an unreadable artifact is rebuilt before it is ever marked reusable. +#[tokio::test] +async fn inline_key_artifact_is_rebuilt_not_reused() { + use std::collections::HashMap; + + use crate::config::MitoConfig; + use crate::series_index::format::SeriesFormats; + use crate::series_index::maintenance::reconcile_series_indexes; + use crate::series_index::version::SeriesIndexVersion; + use crate::test_util::TestEnv; + + let mut env = TestEnv::with_prefix("columnar-samples-inline-keys").await; + let (engine, region) = crate::series_index::tests::prepare_region_with_metadata_and_flush( + &mut env, + &[1000, 2000], + MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + metadata(), + true, + ) + .await; + let format = Arc::new(ColumnarSamplesFormat::default()); + *region.series_index_version_control.formats.write().unwrap() = Arc::new(SeriesFormats { + samples: Some(format.clone()), + ..Default::default() + }); + let store = region.series_index_store.clone().unwrap(); + let (purger, _purges) = series_index_channel(store.clone()); + let reconcile = || { + reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + }; + assert_eq!(1, reconcile().await.unwrap().built_samples); + + // Rewrite every samples manifest to the value an artifact written before the key area carries. + let current = region.series_index_version(); + let mut samples = HashMap::new(); + for (id, handle) in ¤t.samples { + let mut entry = handle.entry().clone(); + for object in &mut entry.artifact.objects { + if !object.path.ends_with("manifest.json") { + continue; + } + let bytes = store.read(&object.path).await.unwrap().to_bytes(); + let mut manifest: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + manifest["identity_encoding"] = serde_json::json!("exact-sparse-key"); + let rewritten = serde_json::to_vec(&manifest).unwrap(); + object.file_size = rewritten.len() as u64; + store.write(&object.path, rewritten).await.unwrap(); + } + entry.artifact.file_size = entry + .artifact + .objects + .iter() + .map(|object| object.file_size) + .sum(); + samples.insert( + *id, + SamplesFileHandle::new(handle.file_id(), entry, purger.clone()), + ); + } + region + .series_index_version_control + .publish(Arc::new(SeriesIndexVersion { + samples: samples.clone(), + range_indexes: current.range_indexes.clone(), + series_indexes: current.series_indexes.clone(), + index_buckets: current.index_buckets.clone(), + })); + + // The manifest alone refuses it: no part has to be read to see the old layout. + let legacy = samples.values().next().unwrap().clone(); + let error = format + .reader(store.clone(), legacy.clone()) + .await + .err() + .expect("an inline-key artifact must be refused at open"); + let message = error.to_string(); + assert!(message.contains("rebuild"), "{message}"); + + // Maintenance must rebuild it, not mark it reusable and fail on the first read. + let stats = reconcile().await.unwrap(); + assert_eq!( + 1, stats.built_samples, + "an artifact the format refuses must be rebuilt" + ); + let rebuilt = region.series_index_version(); + assert_eq!(1, rebuilt.samples.len()); + for handle in rebuilt.samples.values() { + assert_ne!( + handle.entry().artifact.generation, + legacy.entry().artifact.generation, + "the refused artifact must be replaced" + ); + let manifest = manifest(&store, handle).await; + assert_eq!( + IdentityEncoding::CURRENT_SAMPLES, + manifest.identity_encoding + ); + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let batches = reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut times = collect_rows(&batches) + .into_iter() + .map(|row| row.2) + .collect::>(); + times.sort_unstable(); + assert_eq!(vec![1000, 2000], times); + } + use store_api::region_engine::RegionEngine; + engine.stop().await.unwrap(); +} + +/// Fields the frozen v1 manifest still has; unknown-to-the-current-format ones are +/// deliberately not parsed, so the artifact reaches the reader's rejection path. +#[derive(serde::Deserialize)] +struct FrozenV1Manifest { + region_id: u64, + source_file_id: FileId, + generation: FileId, +} + +async fn frozen_v1_handle(store: &ObjectStore) -> SamplesFileHandle { let bytes = include_str!("../../../../../tests/compatibility/wire/mito_columnar_samples/v1-part.hex") .split_whitespace() @@ -857,86 +1351,148 @@ async fn fixture(store: &ObjectStore) -> SamplesFileHandle { let json = include_str!( "../../../../../tests/compatibility/wire/mito_columnar_samples/v1-manifest.json" ); - let manifest: Manifest = serde_json::from_str(json).unwrap(); - store - .write(&manifest.part_path(0), bytes.clone()) - .await - .unwrap(); - store.write(&manifest.path(), json).await.unwrap(); + let frozen: FrozenV1Manifest = serde_json::from_str(json).unwrap(); + let prefix = crate::series_index::samples::samples_prefix(frozen.region_id, frozen.generation); + let manifest_path = format!("{prefix}manifest.json"); + let part_path = format!("{prefix}part-0.samples"); + store.write(&part_path, bytes.clone()).await.unwrap(); + store.write(&manifest_path, json).await.unwrap(); let objects = vec![ crate::series_index::artifact::ArtifactObject { - path: manifest.path(), + path: manifest_path, file_size: json.len() as u64, }, crate::series_index::artifact::ArtifactObject { - path: manifest.part_path(0), + path: part_path, file_size: bytes.len() as u64, }, ]; let artifact = crate::series_index::artifact::ArtifactMetadata { format: ArtifactFormat::ColumnarSamples, - generation: manifest.generation, + generation: frozen.generation, file_size: (json.len() + bytes.len()) as u64, objects, }; let (purger, _) = series_index_channel(store.clone()); SamplesFileHandle::new( RegionFileId::new( - store_api::storage::RegionId::from_u64(manifest.region_id), - manifest.source_file_id, + store_api::storage::RegionId::from_u64(frozen.region_id), + frozen.source_file_id, ), SamplesEntry { source_file_ids: Vec::new(), window: None, source_windows: Default::default(), - file_id: manifest.source_file_id, + file_id: frozen.source_file_id, artifact, }, purger, ) } + +/// A frozen artifact written under hash-only identity must be rejected, never read: +/// it is structurally valid and ordered, so nothing else distinguishes it, and the +/// labels a collapsed block lost cannot be recovered from it. #[tokio::test] -async fn frozen_v1_samples_fixture_and_unknown_codec() { +async fn frozen_v1_samples_fixture_is_rejected() { let store = store(); - let handle = fixture(&store).await; + let handle = frozen_v1_handle(&store).await; let format = ColumnarSamplesFormat::default(); - let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); - let batches = reader - .read(full_request()) + let error = format + .reader(store.clone(), handle) + .await + .err() + .expect("a v1 samples artifact must be rejected"); + // The compatibility property, not the wording of the discriminator: the artifact is + // refused, and the remedy the caller is told about is a rebuild. + let message = error.to_string(); + assert!(message.contains("rebuild"), "{message}"); +} + +/// Rewrites the stored root page so every field column carries an unknown codec, and returns a +/// handle whose descriptor sizes match the rewritten bytes. Everything else stays valid. +async fn unknown_field_codec(store: &ObjectStore, handle: SamplesFileHandle) -> SamplesFileHandle { + let (mut manifest, mut trailer, page, mut bytes) = root_page(store, &handle).await; + let Page::Leaf(mut blocks) = page else { + panic!("the root page must be a leaf") + }; + let mut patched = 0; + for column in blocks.iter_mut().flat_map(|block| block.columns.iter_mut()) { + if column.role == 4 { + column.payload.encoding = 99; + patched += 1; + } + } + assert_eq!(1, patched, "the fixture must store one field column"); + let start = trailer.root.payload.offset as usize; + let encoded = Page::Leaf(blocks).encode(); + bytes.truncate(start); + bytes.extend_from_slice(&encoded); + trailer.root.payload = PayloadReference { + offset: start as u64, + stored: encoded.len() as u32, + decoded: encoded.len() as u32, + checksum: codec::checksum(&encoded), + encoding: 1, + compression: 0, + }; + // The payload area ends where the pages begin, so only the page's own reference and the + // file length change. + trailer.file_length = (bytes.len() + TRAILER_SIZE as usize) as u64; + bytes.extend_from_slice(&trailer.encode()); + let part_path = manifest.part_path(0); + store.write(&part_path, bytes.clone()).await.unwrap(); + manifest.files[0].file_size = bytes.len() as u64; + let manifest_path = manifest.path(); + let json = serde_json::to_vec(&manifest).unwrap(); + store.write(&manifest_path, json.clone()).await.unwrap(); + let mut entry = handle.entry().clone(); + for object in &mut entry.artifact.objects { + if object.path == part_path { + object.file_size = bytes.len() as u64; + } else if object.path == manifest_path { + object.file_size = json.len() as u64; + } else { + panic!("unexpected samples object {}", object.path); + } + } + entry.artifact.file_size = entry.artifact.objects.iter().map(|o| o.file_size).sum(); + let (purger, _receiver) = series_index_channel(store.clone()); + SamplesFileHandle::new(handle.file_id(), entry, purger) +} + +/// A column whose codec the format does not know can only be skipped: a query that does not +/// project it reads normally, and projecting it is an error rather than silent garbage. +#[tokio::test] +async fn unknown_field_codec_can_only_be_skipped() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let handle = build( + &format, + &store, + m.clone(), + &[batch(&m, &[(1, 1, 100, 1, 1, Some(7))], false)], + ) + .await; + let handle = unknown_field_codec(&store, handle).await; + let reader = format.reader(store, handle).await.unwrap(); + let timestamps = reader + .read(request(&["ts"])) .unwrap() .try_collect::>() .await .unwrap(); - assert_eq!( - collect_rows(&batches), - vec![(7, 42, 100, 9, 1, Some(1 << 63)), (7, 42, 100, 8, 0, None)] - ); - let manifest = manifest(&store, &handle).await; - let path = manifest.part_path(0); - let mut bytes = store.read(&path).await.unwrap().to_vec(); - let end = bytes.len() - 112; - let mut trailer = Trailer::decode(&bytes[end..], bytes.len() as u64).unwrap(); - let leaf_start = trailer.root.payload.offset as usize; - let encoding = leaf_start + 12 + 44 + 3 * 44 + 8 + 24; - bytes[encoding..encoding + 2].copy_from_slice(&99u16.to_le_bytes()); - trailer.root.payload.checksum = codec::checksum(&bytes[leaf_start..end]); - bytes[end..].copy_from_slice(&trailer.encode()); - store.write(&path, bytes).await.unwrap(); + assert_eq!(1, timestamps.iter().map(|b| b.num_rows()).sum::()); + let error = reader + .read(request(&["field_0"])) + .unwrap() + .try_collect::>() + .await + .unwrap_err(); assert!( - reader - .read(request(&["ts"])) - .unwrap() - .try_collect::>() - .await - .is_ok() - ); - assert!( - reader - .read(request(&["field_0"])) - .unwrap() - .try_collect::>() - .await - .is_err() + error.to_string().contains("samples column encoding"), + "{error}" ); } @@ -1182,8 +1738,9 @@ async fn batches_column_io_reuses_readers_and_releases_reservations() { assert_eq!(1, creations.load(Ordering::Relaxed)); { let reads = reads.lock().unwrap(); - // Each stream reads framing/root plus one payload read across both blocks. - assert_eq!(8, reads.len()); + // Each stream reads framing/root, one payload read across both blocks, and the part's + // exact keys from the key area. + assert_eq!(10, reads.len()); assert_eq!( 2 * part_size, reads @@ -1199,7 +1756,7 @@ async fn batches_column_io_reuses_readers_and_releases_reservations() { .try_collect::>() .await .unwrap(); - assert_eq!(4, sparse.iter().map(RecordBatch::num_rows).sum::()); + assert_eq!(4, sparse.iter().map(|b| b.num_rows()).sum::()); assert!( reads .lock() @@ -1288,11 +1845,19 @@ async fn boundary_blocks_fetch_remaining_columns_only_for_selected_rows() { .unwrap(); assert!(batches.is_empty()); assert_eq!(1, creations.load(Ordering::Relaxed)); + // Anchor the timestamp read by provenance rather than by position: the timestamp chunk is + // the first chunk of a part, so its read starts at the payload header. Positional indexing + // would silently point at the key-area read instead. let timestamp_read = { let reads = reads.lock().unwrap(); - // Header, trailer, root, then timestamps only. The second part is pruned. - assert_eq!(4, reads.len()); - reads[3].clone() + // Header, trailer, root, the selected block's key, then timestamps only. The other part + // is pruned. + assert_eq!(5, reads.len()); + reads + .iter() + .find(|(_, range)| range.offset() == HEADER_SIZE) + .cloned() + .expect("the timestamp chunk must be read") }; reads.lock().unwrap().clear(); @@ -1308,9 +1873,15 @@ async fn boundary_blocks_fetch_remaining_columns_only_for_selected_rows() { assert_eq!(input[1..5], collect_rows(&batches)); assert_eq!(2, creations.load(Ordering::Relaxed)); let reads = reads.lock().unwrap(); - // Both boundary blocks read timestamps before fetching remaining columns. - assert_eq!(10, reads.len()); - assert_eq!(timestamp_read, reads[3]); + // Both boundary blocks read timestamps before fetching remaining columns, and each part + // reads its own keys. + assert_eq!(12, reads.len()); + assert_eq!( + Some(×tamp_read), + reads + .iter() + .find(|(_, range)| range.offset() == HEADER_SIZE) + ); } #[tokio::test] @@ -1352,28 +1923,38 @@ async fn series_read_windows_are_bounded_and_gap_rows_are_not_returned() { assert_eq!(input, collect_rows(&all)); let (full_bytes, first_window_bytes) = { let reads = reads.lock().unwrap(); - // Header, trailer, one leaf, then three windows of at most 128 blocks. - assert_eq!(6, reads.len()); + // Header, trailer, root (one leaf: a fixed 56-byte descriptor no longer carries the + // series key, so this fixture's 260 blocks fit one page where a key-carrying descriptor + // needed two), three windows of at most 128 blocks, and one merged key read for the + // page's selected keys. + assert_eq!(7, reads.len()); + // The initial read-ahead window is the first read of block 0's first column: it starts + // at the payload header, while the header, trailer, and page reads do not. Anchor on + // that range rather than on an index so a reordered traversal cannot point the + // assertions below at some later window. + let (_, window) = reads + .iter() + .find(|(_, range)| range.offset() == HEADER_SIZE) + .expect("the initial payload window must be read"); + let window_size = window.size().unwrap() as usize; + // One row per block and four payload columns: 8 + 8 + 1 + 8 bytes each. A read that + // spans the whole 128-block window is far larger than a single block's columns. + assert!( + window_size >= 128 * (8 + 8 + 1 + 8), + "the anchored read must be the initial 128-block window, not one block: {window_size}" + ); ( reads.iter().map(|(_, r)| r.size().unwrap()).sum::(), - reads[3].1.size().unwrap() as usize, + window_size, ) }; reads.lock().unwrap().clear(); let mut request = full_request(); - request.ids = Some( - [ - MetricSeriesId { - table_id: 1, - tsid: 0, - }, - MetricSeriesId { - table_id: 1, - tsid: 259, - }, - ] - .into_iter() - .collect(), + request.keys = Some( + [0, 259] + .into_iter() + .map(|tsid| exact(&m, &["a", "b"], 1, tsid)) + .collect(), ); let selected = reader .read(request) @@ -1394,19 +1975,11 @@ async fn series_read_windows_are_bounded_and_gap_rows_are_not_returned() { reads.lock().unwrap().clear(); let mut request = full_request(); - request.ids = Some( - [ - MetricSeriesId { - table_id: 1, - tsid: 0, - }, - MetricSeriesId { - table_id: 1, - tsid: 2, - }, - ] - .into_iter() - .collect(), + request.keys = Some( + [0, 2] + .into_iter() + .map(|tsid| exact(&m, &["a", "b"], 1, tsid)) + .collect(), ); let selected = reader .read(request) @@ -1415,13 +1988,20 @@ async fn series_read_windows_are_bounded_and_gap_rows_are_not_returned() { .await .unwrap(); assert_eq!(vec![input[0], input[2]], collect_rows(&selected)); - // The intervening series is read as a gap, but never decoded or returned. - assert_eq!(4, reads.lock().unwrap().len()); + // The intervening series is read as a gap, but never decoded or returned: only the + // leaf that owns both selected blocks and one payload window are fetched. + assert_eq!(5, reads.lock().unwrap().len()); use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; - // The first pool rejects read-ahead outright. The second admits the first - // 128-block, four-column window but leaves no room for block decoding. - for capacity in [64 * 1024, 2 * first_window_bytes + 128 * 4 * 256] { + // The first pool rejects read-ahead outright. The second admits the first 128-block, + // four-column window but leaves no room for block decoding. Both also hold the part's + // retained exact keys, which are charged to the scan pool before the window is attempted: + // 260 series x (key bytes + the struct + allocator rounding). + let keys_bytes = 260 * (new_sparse_primary_key(&["a", "b"], &m, 1, 0).len() + 32 + 64); + for capacity in [ + 64 * 1024, + 2 * first_window_bytes + 128 * 4 * 256 + keys_bytes, + ] { reads.lock().unwrap().clear(); let pool = Arc::new(GreedyMemoryPool::new(capacity)); let mut request = full_request(); @@ -1437,10 +2017,17 @@ async fn series_read_windows_are_bounded_and_gap_rows_are_not_returned() { { let reads = reads.lock().unwrap(); assert!(reads.len() > 6); + // The fallback reads that same first column one block at a time, so only the larger + // pool's first payload read spans the whole window. Read order depends on the pool, + // which is why the anchor is the range rather than an index. + let initial = reads + .iter() + .find(|(_, range)| range.offset() == HEADER_SIZE) + .map(|(_, range)| range.size().unwrap() as usize); assert_eq!( capacity != 64 * 1024, - reads[3].1.size().unwrap() as usize == first_window_bytes, - "only the larger pool should admit the initial prefetch" + initial == Some(first_window_bytes), + "only the larger pool should admit the initial 128-block prefetch: {initial:?}" ); } let mut stream = reader.read(request).unwrap(); @@ -1467,6 +2054,9 @@ async fn aligned_payload_cache_reuses_gap_bytes_across_readers() { &[batch(&m, &input, false)], ) .await; + // The part's key area, so the loop below can tell a key read from a payload read. + let (_, framing, _, _) = root_page(&store, &handle).await; + let key_area = framing.keys(); let reads = Arc::new(Mutex::new(Vec::new())); let observed = reads.clone(); let store = store.layer( @@ -1499,13 +2089,7 @@ async fn aligned_payload_cache_reuses_gap_bytes_across_readers() { .unwrap() .with_payload_cache(strategy.columnar_payload(namespace)); let mut request = full_request(); - request.ids = Some( - [MetricSeriesId { - table_id: 1, - tsid: id, - }] - .into(), - ); + request.keys = Some([exact(&m, &["a", "b"], 1, id)].into()); request.start = Some(1); request.end = Some(3); let rows = reader @@ -1523,8 +2107,16 @@ async fn aligned_payload_cache_reuses_gap_bytes_across_readers() { collect_rows(&rows) ); if id == 2 { + // A different series reuses the aligned over-read for its payload, including boundary + // values. Its exact key is new, and the key area is part of the same part file, so + // the only read left is the aligned block that covers it. + let reads = reads.lock().unwrap(); assert!( - reads.lock().unwrap().is_empty(), + reads.len() <= 1 + && reads.iter().all(|(_, range)| { + range.offset() >= key_area.start + && range.offset() + range.size().unwrap() <= key_area.end + }), "a different series must reuse the aligned over-read, including boundary values" ); } @@ -1600,11 +2192,153 @@ async fn aligned_payload_cache_assembles_columns_across_blocks() { ); } +/// Keys live in the part's key area, not in the metadata pages: a page keeps a fixed-size +/// reference whatever the key length is, and the area holds one copy per covered series. +#[tokio::test] +async fn keys_are_kept_out_of_metadata_pages() { + let store = store(); + let meta = metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(1024); + let keys = (0..3) + .map(|tsid| new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, tsid)) + .collect::>(); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &keys, &[100, 100, 100], false)], + ) + .await; + let (_, trailer, page, _) = root_page(&store, &handle).await; + assert!( + page.retained_size() < 3 * 1024, + "retained size {} still stores the keys in the page", + page.retained_size() + ); + assert_eq!( + keys.iter() + .map(|key| key.len() as u64 + KEY_CHECKSUM_BYTES) + .sum::(), + trailer.key_end - trailer.payload_end, + "the key area holds one copy of each key and its checksum" + ); +} + +/// A block belongs to exactly one exact series, so its key is emitted once per batch. The +/// scan reservation therefore budgets one key per block instead of one per row: a block of +/// 8192 rows with a 1 KiB key costs 1 KiB of identity, not 8 MiB. +#[tokio::test] +async fn exact_key_is_emitted_once_per_batch() { + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(1024); + let keys = vec![new_sparse_primary_key(&[label.as_str(), "z"], &m, 1, 7); 8192]; + let times = (0..8192_i64).collect::>(); + let handle = build( + &format, + &store, + m.clone(), + &[keyed_batch(&m, &keys, ×, false)], + ) + .await; + let reader = format.reader(store, handle).await.unwrap(); + let mut request = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + request.keys = Some(BTreeSet::from([exact(&m, &[label.as_str(), "z"], 1, 7)])); + // A pooled read that only has to hold one key per batch succeeds; a per-row key column of + // 8 MiB would not fit here. + let pool = Arc::new(GreedyMemoryPool::new(4 * 1024 * 1024)); + request.memory_pool = Some(pool.clone()); + let batches = reader + .read(request) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(8192, batches.iter().map(|b| b.num_rows()).sum::()); + let expected = exact(&m, &[label.as_str(), "z"], 1, 7); + assert!(batches.iter().all(|batch| batch.key == expected)); + // The key is the batch's, not a column: the artifact schema keeps it, output batches do not. + assert!( + batches + .iter() + .all(|batch| batch.batch.column_by_name(PRIMARY_KEY).is_none()) + ); + assert_eq!(0, pool.reserved()); +} + +/// The reservation must cover what a batch actually retains, and identity must not be charged +/// per row: 2135 rows under a 1 KiB key would be over 2 MiB of keys alone. +#[tokio::test] +async fn exact_key_output_reservation_covers_the_retained_array() { + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(1024); + let rows = 2135; + let keys = vec![new_sparse_primary_key(&[label.as_str(), "z"], &m, 1, 7); rows]; + let key_len = keys[0].len(); + let times = (0..rows as i64).collect::>(); + let handle = build( + &format, + &store, + m.clone(), + &[keyed_batch(&m, &keys, ×, false)], + ) + .await; + let reader = format.reader(store.clone(), handle).await.unwrap(); + // The key is returned with the batch, so projecting it as a column is refused instead of + // silently repeating it once per row. + assert!( + reader + .read(SamplesReadRequest { + projection: vec![PRIMARY_KEY.to_string()], + ..Default::default() + }) + .is_err() + ); + let pool = Arc::new(GreedyMemoryPool::new(128 * 1024 * 1024)); + let mut request = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + request.keys = Some(BTreeSet::from([exact(&m, &[label.as_str(), "z"], 1, 7)])); + request.memory_pool = Some(pool.clone()); + let mut stream = reader.read(request).unwrap(); + let batch = stream.try_next().await.unwrap().expect("one batch"); + assert_eq!(rows, batch.num_rows()); + // The batch is yielded while its reservation is still held, so the two are comparable. + let retained = batch.get_array_memory_size(); + let reserved = pool.reserved(); + assert!( + reserved >= retained, + "reservation {reserved} is below the {retained} bytes the yielded batch retains" + ); + // No per-row identity term survives: the estimate is far below one key per row, and the + // key itself is retained once rather than once per row. + assert!( + reserved < rows * key_len, + "reservation {reserved} still charges {rows} keys of {key_len} bytes" + ); + assert!( + reserved < retained + rows * key_len, + "reservation {reserved} dwarfs the {retained} bytes retained" + ); + drop(batch); + drop(stream); + assert_eq!(0, pool.reserved()); +} + /// Sparse selections must bound aligned cache reads as well as selected column bytes. #[rstest::rstest] #[case(1024 * 1024)] #[case(128 * 1024)] -#[case(16 * 1024)] +// Retained exact keys are charged to the scan pool: 128 selected series cost +// 128 * (key bytes + ExactSeriesKey + allocator rounding) ~= 14.3 KB, so a 16 KiB pool can no +// longer hold both the keys and one block's payload reservation. +#[case(64 * 1024)] #[tokio::test] async fn sparse_cached_reads_fit_the_scan_pool(#[case] capacity: usize) { use crate::cache::{CacheManager, CacheStrategy}; @@ -1648,10 +2382,10 @@ async fn sparse_cached_reads_fit_the_scan_pool(#[case] capacity: usize) { let pool = Arc::new(GreedyMemoryPool::new(capacity)); let mut request = full_request(); request.memory_pool = Some(pool.clone()); - request.ids = Some( + request.keys = Some( (0..256) .step_by(2) - .map(|tsid| MetricSeriesId { table_id: 1, tsid }) + .map(|tsid| exact(&m, &["a", "b"], 1, tsid)) .collect(), ); for warm in [false, true] { @@ -1683,3 +2417,5 @@ async fn sparse_cached_reads_fit_the_scan_pool(#[case] capacity: usize) { } } } + +mod keys; diff --git a/src/mito2/src/series_index/samples/tests/keys.rs b/src/mito2/src/series_index/samples/tests/keys.rs new file mode 100644 index 00000000000..36412c33204 --- /dev/null +++ b/src/mito2/src/series_index/samples/tests/keys.rs @@ -0,0 +1,1690 @@ +// 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. + +//! Exact series keys stored outside every bounded page. +//! +//! A block descriptor names its series with a fixed-size reference into the part's key area, +//! so the descriptor — and therefore the metadata page that holds it — never carries a key. +//! These tests pin both directions of that choice: keys far longer than a page must round-trip +//! (they used to exceed the leaf limit and fail the build), and a corrupt reference must be +//! rejected before any row is emitted. + +use std::collections::BTreeSet; +use std::sync::{Arc, Mutex}; + +use datatypes::arrow::array::{Array, RecordBatch, UInt32Array, UInt64Array}; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use futures::TryStreamExt; +use object_store::ObjectStore; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; + +use crate::series_index::artifact::SamplesFileHandle; +use crate::series_index::format::{SamplesBatch, SamplesFormat, SamplesReadRequest}; +use crate::series_index::identity::ExactSeriesKey; +use crate::series_index::purger::series_index_channel; +use crate::series_index::samples::codec; +use crate::series_index::samples::metadata::{ + self, KeyRef, MAGIC, Page, PayloadReference, TRAILER_SIZE, Trailer, +}; +use crate::series_index::samples::{ColumnarSamplesFormat, Manifest}; +use crate::series_index::{TABLE_ID_COLUMN, TSID_COLUMN}; +use crate::test_util::sst_util::new_sparse_primary_key; + +use super::{RecordingRead, build, exact, keyed_batch, request, root_page, sample_batch, store}; + +/// `(table_id, tsid, timestamp)` of every returned row, in output order. +fn collect_times(batches: &[SamplesBatch]) -> Vec<(u32, u64, i64)> { + batches + .iter() + .flat_map(|batch| { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let tsids = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let (times, _) = + datatypes::timestamp::timestamp_array_to_primitive(batch.column(2)).unwrap(); + (0..batch.num_rows()) + .map(|i| (ids.value(i), tsids.value(i), times.value(i))) + .collect::>() + }) + .collect() +} + +fn exact_keys(batches: &[SamplesBatch]) -> Vec { + batches.iter().map(|batch| batch.key.clone()).collect() +} + +/// A region with the 1,024 fields the manifest allows, so the widest possible descriptor and +/// its old, key-length-dependent page budget are exercised. +fn wide_metadata() -> RegionMetadataRef { + let base = super::metadata(); + let mut builder = RegionMetadataBuilder::from_existing((*base).clone()); + for index in 1..1024 { + builder.push_column_metadata(ColumnMetadata { + column_id: 100 + index, + semantic_type: api::v1::SemanticType::Field, + column_schema: ColumnSchema::new( + format!("field_{index}"), + ConcreteDataType::float64_datatype(), + true, + ), + }); + } + Arc::new(builder.build().unwrap()) +} + +/// One block whose series has the given key and every one of the 1,024 fields present. +fn wide_batch(meta: &RegionMetadataRef, key: &[u8], times: &[i64]) -> RecordBatch { + let base = keyed_batch(meta, &vec![key.to_vec(); times.len()], times, false); + let mut fields = base.schema().fields().to_vec(); + let mut arrays = base.columns().to_vec(); + for index in 1..1024 { + fields.insert( + index, + Arc::new(Field::new( + format!("field_{index}"), + DataType::Float64, + true, + )), + ); + arrays.insert(index, base.column(0).clone()); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap() +} + +/// Replaces the part and manifest of `handle` with `bytes` and their recorded sizes. +/// +/// Everything else stays valid, so a test can isolate one rewritten descriptor. +async fn rewrite_part( + store: &ObjectStore, + handle: SamplesFileHandle, + mut manifest: Manifest, + bytes: Vec, +) -> SamplesFileHandle { + let part_path = manifest.part_path(0); + store.write(&part_path, bytes.clone()).await.unwrap(); + manifest.files[0].file_size = bytes.len() as u64; + let manifest_path = manifest.path(); + let json = serde_json::to_vec(&manifest).unwrap(); + store.write(&manifest_path, json.clone()).await.unwrap(); + let mut entry = handle.entry().clone(); + for object in &mut entry.artifact.objects { + if object.path == part_path { + object.file_size = bytes.len() as u64; + } else if object.path == manifest_path { + object.file_size = json.len() as u64; + } else { + panic!("unexpected samples object {}", object.path); + } + } + entry.artifact.file_size = entry.artifact.objects.iter().map(|o| o.file_size).sum(); + let (purger, _receiver) = series_index_channel(store.clone()); + SamplesFileHandle::new(handle.file_id(), entry, purger) +} + +/// A single label this long encodes to a key well past the 65,536-byte page limit. The old +/// descriptor had to fit `48 + 44 * C + key_length` bytes plus the leaf prefix in one page, +/// so this build failed with "sample descriptor exceeds leaf limit"; the key now lives in its +/// own area and only a 12-byte reference stays in the descriptor. +#[tokio::test] +async fn key_longer_than_one_page_round_trips() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(65 * 1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 7); + assert!(key.len() > 65_536, "key is {} bytes", key.len()); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[key.clone(), key.clone()], + &[100, 200], + false, + )], + ) + .await; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let selection = exact(&meta, &[label.as_str(), "z"], 1, 7); + assert_eq!(selection.bytes(), key.as_slice()); + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.keys = Some(BTreeSet::from([selection.clone()])); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + // Both rows belong to one exact series, so they come back as one batch carrying that key. + assert_eq!(vec![(1, 7, 100), (1, 7, 200)], collect_times(&batches)); + assert_eq!(vec![selection], exact_keys(&batches)); + // The stored key is the writer's key byte for byte. + let (manifest, trailer, page, bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!(1, blocks.len()); + let reference = blocks[0].key; + assert_eq!(reference.offset, trailer.payload_end); + assert_eq!(reference.length as usize, key.len()); + // One copy is the key plus its stored checksum, which the reference does not count. + assert_eq!( + trailer.key_end - trailer.payload_end, + key.len() as u64 + metadata::KEY_CHECKSUM_BYTES + ); + assert_eq!( + key.as_slice(), + &bytes[reference.offset as usize..reference.end().unwrap() as usize] + ); + assert_eq!( + codec::checksum(key.as_slice()), + u64::from_le_bytes( + bytes[reference.end().unwrap() as usize..reference.checksum_end().unwrap() as usize] + .try_into() + .unwrap() + ) + ); + // A key longer than a page still leaves a descriptor far below the page limit. + assert_eq!( + 56 + 44 * 4, + metadata::descriptor_bytes(blocks[0].columns.len()) + ); + assert!(metadata::descriptor_bytes(blocks[0].columns.len()) + 12 < 65_536 / 2); + // The key area is part of the part, not of any page. + assert!(trailer.key_end > trailer.payload_end); + assert!(manifest.files[0].file_size > trailer.key_end); +} + +/// At the 1,024-field maximum the old layout left `65,536 - (12 + 48 + 44 * 1,027) = 20,288` +/// bytes for the key. This key is longer than that and still round-trips with every field +/// present in the block. +#[tokio::test] +async fn many_fields_round_trip_a_key_past_the_old_ceiling() { + let store = store(); + let meta = wide_metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(20 * 1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 3, 11); + assert!(key.len() > 20_288, "key is {} bytes", key.len()); + let handle = build( + &format, + &store, + meta.clone(), + &[wide_batch(&meta, &key, &[10, 20])], + ) + .await; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let selection = exact(&meta, &[label.as_str(), "z"], 3, 11); + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.keys = Some(BTreeSet::from([selection.clone()])); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(vec![(3, 11, 10), (3, 11, 20)], collect_times(&batches)); + assert_eq!(vec![selection], exact_keys(&batches)); + // The widest descriptor stays a fixed number of bytes whatever the key's length is. + let (_, trailer, page, _) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!( + 56 + 44 * 1027, + metadata::descriptor_bytes(blocks[0].columns.len()) + ); + assert_eq!(3 + 1024, blocks[0].columns.len()); + assert_eq!( + trailer.key_end - trailer.payload_end, + key.len() as u64 + metadata::KEY_CHECKSUM_BYTES + ); +} + +/// A reference that points outside the key area is a corrupt artifact, not a reason to read +/// some other series' bytes: decoding the page it appears in must fail, and a read must fail +/// before it emits any row. +#[tokio::test] +async fn corrupt_key_reference_is_rejected_before_rows() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "y".repeat(8 * 1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 9); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + std::slice::from_ref(&key), + &[100], + false, + )], + ) + .await; + let (manifest, trailer, page, bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + let area = trailer.keys(); + // Past the end of the area, before its start, and a zero length: each one addresses + // descriptor bytes that cannot be a key of this part. + let corrupt = [ + KeyRef { + offset: area.end - 1, + length: 4, + }, + KeyRef { + offset: area.start - 4, + length: 4, + }, + KeyRef { + offset: area.start, + length: 0, + }, + ]; + let start = trailer.root.payload.offset as usize; + let mut encoded = Page::Leaf(blocks.clone()).encode(); + for bad in corrupt { + let mut patched = blocks.clone(); + patched[0].key = bad; + encoded = Page::Leaf(patched).encode(); + assert!( + Page::decode( + &encoded, + &trailer.root, + trailer.depth, + trailer.areas(), + &manifest + ) + .is_err(), + "{bad:?} must not decode" + ); + } + // Install the last patch and read it back through the reader. + let mut rewritten = bytes; + rewritten.truncate(start); + rewritten.extend_from_slice(&encoded); + let mut trailer = trailer; + trailer.root.payload = PayloadReference { + offset: start as u64, + stored: encoded.len() as u32, + decoded: encoded.len() as u32, + checksum: codec::checksum(&encoded), + encoding: 1, + compression: 0, + }; + trailer.file_length = (rewritten.len() + TRAILER_SIZE as usize) as u64; + rewritten.extend_from_slice(&trailer.encode()); + let handle = rewrite_part(&store, handle, manifest, rewritten).await; + let reader = format.reader(store, handle).await.unwrap(); + let mut stream = reader.read(request(&[TABLE_ID_COLUMN])).unwrap(); + let error = stream + .try_next() + .await + .expect_err("a corrupt key reference must fail the read"); + assert!( + error.to_string().contains("key area"), + "unexpected error: {error}" + ); +} + +/// Two distinct series that share one `(table_id, tsid)`, one with a long label. Each block +/// still carries its own exact key, so neither read can see the partner's samples, and the +/// key area holds exactly one copy per series. +#[tokio::test] +async fn colliding_series_keep_their_own_exact_keys() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let long_label = format!("host-{}", "z".repeat(300)); + let first = new_sparse_primary_key(&["host-020078", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&[long_label.as_str(), "z"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[ + keyed_batch(&meta, &[first.clone(), first.clone()], &[100, 200], false), + keyed_batch( + &meta, + &[partner.clone(), partner.clone()], + &[100, 200], + false, + ), + ], + ) + .await; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let first_key = exact(&meta, &["host-020078", "z"], 1, 7); + let partner_key = exact(&meta, &[long_label.as_str(), "z"], 1, 7); + assert_ne!(first_key, partner_key); + assert_eq!(first_key.id(), partner_key.id()); + for (selection, expected) in [ + (first_key.clone(), vec![(1, 7, 100), (1, 7, 200)]), + (partner_key.clone(), vec![(1, 7, 100), (1, 7, 200)]), + ] { + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.keys = Some(BTreeSet::from([selection.clone()])); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(expected, collect_times(&batches)); + assert_eq!(vec![selection], exact_keys(&batches)); + } + // One copy per series: two references, in block order, and the area holds both keys. + let (_, trailer, page, bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!(2, blocks.len()); + assert_eq!(first_key.id(), blocks[0].id); + assert_eq!(partner_key.id(), blocks[0].id); + assert_ne!(blocks[0].key, blocks[1].key); + assert!(blocks[0].key.end().unwrap() <= blocks[1].key.offset); + assert_eq!(trailer.payload_end, blocks[0].key.offset); + assert_eq!( + first.len() + partner.len() + 2 * metadata::KEY_CHECKSUM_BYTES as usize, + (trailer.key_end - trailer.payload_end) as usize + ); + assert_eq!( + first.as_slice(), + &bytes[blocks[0].key.offset as usize..blocks[0].key.end().unwrap() as usize] + ); + assert_eq!( + partner.as_slice(), + &bytes[blocks[1].key.offset as usize..blocks[1].key.end().unwrap() as usize] + ); +} + +/// A part written before the key area has a 112-byte trailer and no `key_end`. The reader locates +/// a trailer from the recorded object size, so it always reads the last 120 bytes: the shifted +/// frame must be rejected, never decoded as if it were current. There is no version bump to tell +/// the two shapes apart, so an artifact from an earlier build of this branch is unreadable and +/// has to be rebuilt. +#[tokio::test] +async fn legacy_112_byte_trailer_is_rejected() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let key = new_sparse_primary_key(&["a", "b"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &[key], &[100], false)], + ) + .await; + let (_, trailer, _, _) = root_page(&store, &handle).await; + assert_eq!(120, TRAILER_SIZE); + // The pre-key-area frame: the same fields without `key_end`, which is 112 bytes. + let mut legacy = trailer.encode(); + // Drop `key_end` (96..104) and the current magic/checksum, then seal the old way: the old + // checksum covers its 96 body bytes plus the magic, 104 in total. + legacy.drain(96..104); + legacy.truncate(96); + legacy.extend_from_slice(MAGIC); + let sum = codec::checksum(&legacy); + legacy.extend_from_slice(&sum.to_le_bytes()); + assert_eq!(112, legacy.len()); + // A 112-byte slice is not a trailer of this format at all. + let error = Trailer::decode(&legacy, trailer.file_length) + .err() + .expect("a 112-byte frame is not a trailer"); + assert!(error.to_string().contains("truncated"), "{error}"); + // What a reader actually does with an old part: the last 120 bytes of the file, which start + // 8 bytes before the old trailer. The magic still lands where this format expects it, so the + // frame is caught by the checksum over the shifted prefix rather than by the magic. + let mut window: Vec = vec![0; 8]; + window.extend_from_slice(&legacy); + assert_eq!(TRAILER_SIZE as usize, window.len()); + assert!(Trailer::decode(&window, trailer.file_length).is_err()); +} + +/// One request's key-area traffic: `(reads, bytes)` plus the rows it returned. +async fn key_area_traffic( + store: &ObjectStore, + format: &ColumnarSamplesFormat, + handle: &SamplesFileHandle, + query: SamplesReadRequest, +) -> (usize, u64, usize) { + use object_store::layers::mock::{MockLayerBuilder, oio}; + + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner: oio::Reader| { + Box::new(RecordingRead { + inner, + path: path.into(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(); + let reader = format + .reader(store.clone().layer(layer), handle.clone()) + .await + .unwrap(); + // The manifest was read above; only the query's own reads matter here. + reads.lock().unwrap().clear(); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let rows = batches.iter().map(|batch| batch.num_rows()).sum::(); + let (_, trailer, _, _) = root_page(store, handle).await; + let area = trailer.keys(); + let reads = reads.lock().unwrap(); + let (mut count, mut bytes) = (0usize, 0u64); + for (_, range) in reads.iter() { + let start = range.offset(); + let end = start + range.size().unwrap(); + if start < end && start >= area.start && end <= area.end { + count += 1; + bytes += end - start; + } + } + (count, bytes, rows) +} + +/// The reader stays lazy and does not pay one key read per series. +/// +/// 260 one-row series share one part, so a full scan touches 260 distinct keys. Reading each of +/// them from the key area costs 260 tiny reads next to the three coalesced payload windows the +/// same scan issues; merging the selected blocks of a page into bounded reads keeps the traffic +/// at a handful of reads while still fetching exactly the bytes of the selected keys. +#[tokio::test] +async fn key_reads_are_merged_per_page() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let input = (0..260u64) + .map(|id| (1, id, 0, 1, 1, Some(id))) + .collect::>(); + let handle = build( + &format, + &store, + meta.clone(), + &[super::batch(&meta, &input, false)], + ) + .await; + let key_len = new_sparse_primary_key(&["a", "b"], &meta, 1, 0).len() as u64 + + metadata::KEY_CHECKSUM_BYTES; + let (reads, bytes, rows) = + key_area_traffic(&store, &format, &handle, request(&[TABLE_ID_COLUMN])).await; + assert_eq!(260, rows); + // Every selected key is fetched exactly once: no duplicate bytes, no whole-area read. + assert_eq!(260 * key_len, bytes, "key bytes fetched"); + // One read for the whole page: before the keys of a page were merged, this scan issued one + // key-area read per series (260 of them, about 30 bytes each). + assert_eq!( + 1, reads, + "{reads} key-area reads for 260 series is per-series amplification" + ); + + // Selection stays lazy: two series cost a bounded read, not the whole area. + let mut query = request(&[TABLE_ID_COLUMN]); + query.keys = Some(BTreeSet::from([ + exact(&meta, &["a", "b"], 1, 0), + exact(&meta, &["a", "b"], 1, 259), + ])); + let (reads, bytes, rows) = key_area_traffic(&store, &format, &handle, query).await; + assert_eq!(2, rows); + // The two keys are ~8 KiB apart, past the 4 KiB gap the reader is allowed to over-read, so + // they stay two bounded reads rather than one whole-area read. + assert_eq!(2, reads, "{reads} key-area reads for two selected series"); + assert_eq!( + 2 * key_len, + bytes, + "{bytes} bytes for two keys means the whole key area was read" + ); +} + +/// A key area whose bytes disagree with the descriptor that references them is a corrupt part. +/// +/// The reserved `(table_id, tsid)` prefix is the first 12 bytes of a canonical sparse key, so +/// flipping one byte there keeps the key parseable while it names another series. +#[tokio::test] +async fn key_bytes_that_disagree_with_the_descriptor_are_rejected() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let key = new_sparse_primary_key(&["a", "b"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &[key], &[100], false)], + ) + .await; + let (manifest, trailer, page, mut bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + let reference = blocks[0].key; + assert_eq!(reference.offset, trailer.payload_end); + // The canonical sparse prefix is `[column id 0..4][table_id marker 4][table_id 5..9]` + // `[tsid column id 9..13][tsid marker 13][tsid 14..22]`, then the labels. Byte 16 is in the + // middle of the tsid value, so the framing still parses and only the identifier changes. + bytes[reference.offset as usize + 16] ^= 1; + // Repair the stored checksum: this test is about the key disagreeing with its descriptor, not + // about the bytes being corrupted, so the checksum must not be what rejects the part. + let sum = codec::checksum(&bytes[reference.offset as usize..reference.end().unwrap() as usize]); + let checksum_start = reference.end().unwrap() as usize; + bytes[checksum_start..reference.checksum_end().unwrap() as usize] + .copy_from_slice(&sum.to_le_bytes()); + store.write(&manifest.part_path(0), bytes).await.unwrap(); + let reader = format.reader(store, handle).await.unwrap(); + let mut stream = reader.read(request(&[TABLE_ID_COLUMN])).unwrap(); + let error = stream + .try_next() + .await + .expect_err("a key that disagrees with its descriptor must fail the read"); + assert!( + error.to_string().contains("disagrees with its identifier"), + "unexpected error: {error}" + ); +} + +/// The trailer and root page of one part, plus its bytes. +/// +/// `root_page` insists on a single-part artifact; the part-boundary tests need one part at a time. +async fn part_root( + store: &ObjectStore, + handle: &SamplesFileHandle, + part: u32, +) -> (Trailer, Page, Vec) { + let manifest = super::manifest(store, handle).await; + let bytes = store + .read(&manifest.part_path(part)) + .await + .unwrap() + .to_bytes() + .to_vec(); + let start = bytes.len() - TRAILER_SIZE as usize; + let trailer = Trailer::decode(&bytes[start..], bytes.len() as u64).unwrap(); + assert_eq!(1, trailer.depth, "one leaf keeps the root page addressable"); + let offset = trailer.root.payload.offset as usize; + let stored = &bytes[offset..offset + trailer.root.payload.stored as usize]; + let page = Page::decode( + &trailer.root.payload.unpack(stored).unwrap(), + &trailer.root, + trailer.depth, + trailer.areas(), + &manifest, + ) + .unwrap(); + (trailer, page, bytes) +} + +/// Rewrites the two stored key entries of a two-series part in place. +/// +/// The descriptors stay untouched and the part keeps its length, so whatever rejects the part +/// must be the exact keys themselves, not the framing or the reference order. +async fn rewrite_entries( + store: &ObjectStore, + handle: SamplesFileHandle, + edit: impl FnOnce(&mut [u8], [std::ops::Range; 2]), +) { + let (manifest, _, page, mut bytes) = root_page(store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!(2, blocks.len(), "the fixture holds two series"); + let entries = [&blocks[0].key, &blocks[1].key] + .map(|reference| reference.offset as usize..reference.checksum_end().unwrap() as usize); + edit(&mut bytes, entries); + store.write(&manifest.part_path(0), bytes).await.unwrap(); +} + +/// One query's key-area reads - `(reads, bytes)` - and row count, run through a recording layer. +/// +/// A read counts when it *intersects* the key area, not only when it lies inside it: an aligned +/// read is clipped at the area end but may start before the area, and it still fetches key bytes. +async fn recorded_key_area_reads( + store: &ObjectStore, + format: &ColumnarSamplesFormat, + handle: &SamplesFileHandle, + query: SamplesReadRequest, +) -> (usize, u64, crate::error::Result) { + use object_store::layers::mock::{MockLayerBuilder, oio}; + + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner: oio::Reader| { + Box::new(RecordingRead { + inner, + path: path.into(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(); + let reader = format + .reader(store.clone().layer(layer), handle.clone()) + .await + .unwrap(); + reads.lock().unwrap().clear(); + let result = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .map(|batches| batches.iter().map(|batch| batch.num_rows()).sum::()); + let (_, trailer, _, _) = root_page(store, handle).await; + let area = trailer.keys(); + let reads = reads.lock().unwrap(); + let (mut count, mut bytes) = (0usize, 0u64); + for (_, range) in reads.iter() { + let start = range.offset(); + let end = start + range.size().unwrap(); + if start < end && start < area.end && end > area.start { + count += 1; + bytes += end - start; + } + } + (count, bytes, result) +} + +/// One exact series may span parts, and every part lays out its own key area. +/// +/// The two copies of one series therefore live at unrelated offsets. Comparing references across +/// the boundary reads those offsets as one area, so a legitimate artifact looks overlapping (or +/// like a new series that skipped its ordinal). Only the exact key bytes can tell a series that +/// continues from a series that starts. +#[tokio::test] +async fn same_series_across_parts_with_different_key_offsets_is_readable() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat { + block_rows: 1, + part_bytes: 1, + ..Default::default() + }; + // The second row has no field value, so its block stores one chunk fewer and the part it + // opens places its key area before the previous part did. + let input = vec![ + (1, 7, 10, 1, 1, Some(7)), + (1, 7, 20, 1, 1, None), + (1, 7, 30, 1, 1, Some(9)), + ]; + let handle = build( + &format, + &store, + meta.clone(), + &[super::batch(&meta, &input, false)], + ) + .await; + assert_eq!( + 3, + super::manifest(&store, &handle).await.files.len(), + "every block rolls its part" + ); + let (first, _, _) = part_root(&store, &handle, 0).await; + let (second, _, _) = part_root(&store, &handle, 1).await; + assert_ne!( + first.payload_end, second.payload_end, + "the parts must place their key areas differently" + ); + let reader = format.reader(store, handle).await.unwrap(); + let batches = reader + .read(request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + vec![(1, 7, 10), (1, 7, 20), (1, 7, 30)], + collect_times(&batches) + ); + // One series, so every batch reports the same exact key. + assert_eq!( + vec![exact(&meta, &["a", "b"], 1, 7); 3], + exact_keys(&batches) + ); +} + +/// Two parts' key areas are unrelated, so equal references in them name different series. +/// +/// Both series are alone in their part and encode to the same length, so both key areas start at +/// the same offset and the two references are the same number. Treating an equal reference as one +/// series would then demand a continuation the artifact never promised: the second series starts +/// its own timestamps and ordinal. +#[tokio::test] +async fn equal_key_references_across_parts_are_different_series() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat { + block_rows: 1, + part_bytes: 1, + ..Default::default() + }; + let first = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-000002", "z"], &meta, 1, 7); + assert_eq!(first.len(), partner.len(), "the references must coincide"); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &[first, partner], &[100, 50], false)], + ) + .await; + assert_eq!(2, super::manifest(&store, &handle).await.files.len()); + let (_, first_page, _) = part_root(&store, &handle, 0).await; + let (_, second_page, _) = part_root(&store, &handle, 1).await; + let Page::Leaf(first_blocks) = first_page else { + panic!("one leaf page holds the first part") + }; + let Page::Leaf(second_blocks) = second_page else { + panic!("one leaf page holds the second part") + }; + assert_eq!( + first_blocks[0].key, second_blocks[0].key, + "the two series' references must be the same number" + ); + assert!( + second_blocks[0].min < first_blocks[0].max, + "the second series must restart its timestamps" + ); + let reader = format.reader(store, handle).await.unwrap(); + let batches = reader + .read(request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(vec![(1, 7, 100), (1, 7, 50)], collect_times(&batches)); + assert_eq!( + vec![ + exact(&meta, &["host-000001", "z"], 1, 7), + exact(&meta, &["host-000002", "z"], 1, 7), + ], + exact_keys(&batches) + ); +} + +/// A label byte is not covered by the key's own framing, so only a stored checksum protects it. +/// +/// Flipping one byte of a label value keeps the sparse framing and the reserved `(table_id, tsid)` +/// prefix valid, so the key still parses and still claims the descriptor's identifier. Without a +/// per-key checksum the block's rows would be attributed to another series and returned as a +/// perfectly ordinary batch. +#[tokio::test] +async fn corrupted_label_byte_is_rejected() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let key = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + std::slice::from_ref(&key), + &[100], + false, + )], + ) + .await; + let (manifest, trailer, page, mut bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + let reference = blocks[0].key; + assert_eq!(reference.offset, trailer.payload_end); + // Find a label byte whose flip keeps the key well framed and under the same identifier. + let mut offsets = mito_codec::row_converter::SparseOffsetsCache::new(); + let position = (22..key.len()) + .rev() + .find(|position| { + let mut candidate = key.clone(); + candidate[*position] ^= 1; + ExactSeriesKey::from_sparse_key(&candidate, &mut offsets) + .is_ok_and(|decoded| decoded.id() == blocks[0].id) + }) + .expect("a label byte whose flip keeps the key valid"); + bytes[reference.offset as usize + position] ^= 1; + store.write(&manifest.part_path(0), bytes).await.unwrap(); + let reader = format.reader(store, handle).await.unwrap(); + let error = reader + .read(request(&[TABLE_ID_COLUMN])) + .unwrap() + .try_next() + .await + .expect_err("a rewritten label must not silently name another series"); + assert!( + error.to_string().contains("sample key checksum mismatch"), + "unexpected error: {error}" + ); +} + +/// References can be ordered and disjoint while the keys they point at are reversed. +/// +/// The structural check only sees descriptor fields, so an increasing reference holding a smaller +/// exact key than its predecessor passes it. Two distinct series under one `(table_id, tsid)` are +/// exactly this case once the key area is corrupted: downstream reads rely on the batches arriving +/// in exact-key order, so the loaded keys themselves must be ordered. +#[tokio::test] +async fn reversed_exact_keys_under_ordered_references_are_rejected() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-000002", "z"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[first.clone(), partner.clone()], + &[100, 200], + false, + )], + ) + .await; + // Swapping whole entries keeps both references, both checksums, and both lengths valid: the + // only thing wrong is that the first block is now the larger key. + rewrite_entries(&store, handle.clone(), |bytes, [left, right]| { + let left_entry = bytes[left.clone()].to_vec(); + bytes.copy_within(right.clone(), left.start); + bytes[right].copy_from_slice(&left_entry); + }) + .await; + let reader = format.reader(store, handle).await.unwrap(); + // The stream is lazy, so the second block is only visited after the first is yielded. + let error = reader + .read(request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"])) + .unwrap() + .try_collect::>() + .await + .expect_err("reversed exact keys must not be read as ordered series"); + assert!( + error + .to_string() + .contains("unordered sample exact series keys"), + "unexpected error: {error}" + ); +} + +/// Two references that point at one exact key are one series, so it must continue. +/// +/// A different reference alone used to mean "a new series", which let a second copy of the same +/// key restart its ordinal and its timestamps. The key bytes say otherwise: the same series is +/// one series, and a block that starts at ordinal zero and goes back in time is corrupt. +#[tokio::test] +async fn repeated_exact_key_under_new_reference_must_continue() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let partner = new_sparse_primary_key(&["host-000002", "z"], &meta, 1, 7); + assert_eq!(first.len(), partner.len()); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[first.clone(), partner], + &[100, 50], + false, + )], + ) + .await; + // The second entry becomes a second copy of the first series' key, with a valid checksum. + rewrite_entries(&store, handle.clone(), |bytes, [left, right]| { + let key = bytes[left.start..left.end - metadata::KEY_CHECKSUM_BYTES as usize].to_vec(); + let sum = codec::checksum(&key); + bytes[right.start..right.end - metadata::KEY_CHECKSUM_BYTES as usize].copy_from_slice(&key); + bytes[right.end - metadata::KEY_CHECKSUM_BYTES as usize..right.end] + .copy_from_slice(&sum.to_le_bytes()); + }) + .await; + let reader = format.reader(store, handle).await.unwrap(); + let error = reader + .read(request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"])) + .unwrap() + .try_collect::>() + .await + .expect_err("a repeated key that restarts its series must not be read"); + assert!( + error.to_string().contains("unordered sample blocks"), + "unexpected error: {error}" + ); +} + +/// A pool that cannot pay for one key must fail before the key area is read. +/// +/// The fetched entry and the owned key coexist, so the pool has to hold both before the read. A +/// reader that fetches first and charges afterwards has already issued I/O for bytes it cannot +/// keep, which is what "no I/O once the pool cannot pay" rules out. +#[tokio::test] +async fn unpayable_key_prefetch_issues_no_key_read() { + use common_error::ext::ErrorExt; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let label = "x".repeat(4096); + let keys = (0..32u64) + .map(|tsid| new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, tsid)) + .collect::>(); + let times = vec![100; keys.len()]; + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &keys, ×, false)], + ) + .await; + // Four KiB cannot hold one 4 KiB key entry plus the owned key decoded from it. + let pool = Arc::new(GreedyMemoryPool::new(4 * 1024)); + let mut query = super::full_request(); + query.memory_pool = Some(pool.clone()); + let (reads, _, result) = recorded_key_area_reads(&store, &format, &handle, query).await; + let error = result.expect_err("an unpayable key must fail the read"); + assert_eq!( + common_error::status_code::StatusCode::RuntimeResourcesExhausted, + error.status_code() + ); + assert_eq!( + 0, reads, + "{reads} key-area reads happened although the pool cannot hold one key" + ); + assert_eq!(0, pool.reserved()); +} + +/// A cached part answers a key from the payload cache, never from the key area. +/// +/// Keys use the same aligned blocks as column chunks, so a part that is already cached must not +/// pay a key-area read - including when the scan pool shapes how the read is issued. +/// +/// This is a cache-behaviour guard, not a proof of the key-fetch budget: it passes without the +/// peak reservation too, because a warm part issues no read either way. The budget itself is +/// pinned by `unpayable_key_prefetch_issues_no_key_read` and +/// `per_key_fallback_reserves_the_aligned_peak`. +#[tokio::test] +async fn cached_parts_need_no_key_read() { + use crate::cache::{CacheManager, CacheStrategy}; + use crate::series_index::format::SamplesFormatReader; + use crate::series_index::samples::reader; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use store_api::storage::FileId; + + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let keys = (0..4u64) + .map(|tsid| new_sparse_primary_key(&["a", "b"], &meta, 1, tsid)) + .collect::>(); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &keys, &[100, 100, 100, 100], false)], + ) + .await; + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .page_cache_size(16 * 1024 * 1024) + .build(), + )); + let cache = strategy.columnar_payload(FileId::random()); + let open = |store: &ObjectStore| { + let store = store.clone(); + let handle = handle.clone(); + let cache = cache.clone(); + async move { + reader::SamplesReader::open(store, handle) + .await + .unwrap() + .with_payload_cache(cache) + } + }; + // Warm every aligned block of the part, keys included. + let warm = open(&store).await; + warm.read(super::full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .reader_factory(Arc::new( + move |path, _, inner: object_store::layers::mock::oio::Reader| { + Box::new(RecordingRead { + inner, + path: path.into(), + reads: observed.clone(), + }) + }, + )) + .build() + .unwrap(); + let layered = store.clone().layer(layer); + let reader = open(&layered).await; + reads.lock().unwrap().clear(); + let pool = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut query = super::full_request(); + query.memory_pool = Some(pool.clone()); + // One selected series still needs its key, which this part already has cached. + query.keys = Some(BTreeSet::from([exact(&meta, &["a", "b"], 1, 2)])); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(vec![(1, 2, 100)], collect_times(&batches)); + // Counting any read that intersects the key area, so an aligned read starting before it is + // counted too: a warm part must answer the key from the payload cache either way. + let (_, trailer, _, _) = root_page(&store, &handle).await; + let area = trailer.keys(); + let reads = reads.lock().unwrap(); + let key_reads = reads + .iter() + .filter(|(_, range)| { + let start = range.offset(); + let end = start + range.size().unwrap(); + start < end && start < area.end && end > area.start + }) + .count(); + assert_eq!( + 0, key_reads, + "{key_reads} key-area reads on a part whose blocks are already cached" + ); + assert_eq!(0, pool.reserved()); +} + +/// One writer over one long-keyed series, returning the accounting after the second block flushes +/// plus the size of one key copy. +/// +/// `root_capacity` is one and `page_bytes` admits exactly one maximal descriptor per leaf, so the +/// second block starts a new part while the series continues into it. +async fn cross_rollover_accounting(part_bytes: u64) -> ((u64, u64, u64, usize), u64) { + use crate::series_index::format::FormatWriter; + use crate::series_index::samples::writer::SamplesWriter; + use store_api::storage::FileId; + + let store = store(); + let meta = super::metadata(); + let label = "x".repeat(512); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 7); + let copy_bytes = key.len() as u64 + metadata::KEY_CHECKSUM_BYTES; + let format = ColumnarSamplesFormat { + block_rows: 1, + part_bytes, + page_bytes: 244, + root_capacity: 1, + ..Default::default() + }; + let generation = FileId::random(); + let prefix = format!("{}/samples/{generation}/", meta.region_id.as_u64()); + let mut writer = SamplesWriter::new( + format, + meta.clone(), + store, + &prefix, + FileId::random(), + generation, + ) + .unwrap(); + // A row is buffered before the next one flushes its block, so three single-row batches flush + // two blocks: the second one is the block that starts the new part. + for (time, value) in [(10, None), (20, Some(0.0)), (30, Some(0.0))] { + let batch = sample_batch( + &meta, + std::slice::from_ref(&key), + vec![time], + vec![1], + vec![0], + vec![value], + false, + ); + writer.write(&batch).await.unwrap(); + } + (writer.part_accounting(), copy_bytes) +} + +/// A series that crosses a metadata rollover must have its new part's key copy counted. +/// +/// `finish_part` clears the part's key accounting, so a "is this the first copy" decision taken +/// before the rollover skips the copy the new part makes. The part target then never sees the key +/// bytes and defers the part boundary by one part's worth of keys. +#[tokio::test] +async fn series_across_a_metadata_rollover_counts_its_key_copy() { + let ((position, counted, blocks, parts), copy_bytes) = + cross_rollover_accounting(u64::MAX).await; + assert_eq!(1, parts, "the metadata rule must complete one part"); + assert_eq!(1, blocks, "the second block opens the new part"); + assert_eq!( + copy_bytes, counted, + "the new part must count its own copy of the series key" + ); + // `flush_block` compares exactly `position + part_key_bytes` with the part target, so a part + // target of that sum rolls the new part at once - only when the copy is counted. + let ((_, _, blocks_again, parts_again), _) = + cross_rollover_accounting(position + counted).await; + assert_eq!( + 2, parts_again, + "the counted key copy must be what reaches the part target" + ); + assert_eq!( + 0, blocks_again, + "the part target rolls the new part at once" + ); +} + +/// A memory pool that remembers the highest reservation it ever held and how much each named +/// consumer held, so a test can see that a scan charged for something while it was alive. +#[derive(Debug)] +struct WatchingPool { + inner: datafusion::execution::memory_pool::GreedyMemoryPool, + high: std::sync::atomic::AtomicUsize, + consumers: Mutex>, +} + +impl std::fmt::Display for WatchingPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + std::fmt::Display::fmt(&self.inner, f) + } +} + +impl WatchingPool { + fn new(limit: usize) -> Self { + Self { + inner: datafusion::execution::memory_pool::GreedyMemoryPool::new(limit), + high: std::sync::atomic::AtomicUsize::new(0), + consumers: Mutex::new(std::collections::HashMap::new()), + } + } + + /// The most bytes any scan using this pool ever had reserved at once. + fn high_water(&self) -> usize { + self.high.load(std::sync::atomic::Ordering::Relaxed) + } + + /// Bytes currently reserved by one named consumer. + fn consumer(&self, name: &str) -> usize { + self.consumers + .lock() + .unwrap() + .get(name) + .copied() + .unwrap_or(0) + } + + fn record(&self, name: &str, bytes: usize, grow: bool) { + { + let mut consumers = self.consumers.lock().unwrap(); + let slot = consumers.entry(name.to_string()).or_default(); + if grow { + *slot += bytes; + } else { + *slot = slot.saturating_sub(bytes); + } + } + self.high.fetch_max( + datafusion::execution::memory_pool::MemoryPool::reserved(&self.inner), + std::sync::atomic::Ordering::Relaxed, + ); + } +} + +impl datafusion::execution::memory_pool::MemoryPool for WatchingPool { + fn name(&self) -> &str { + datafusion::execution::memory_pool::MemoryPool::name(&self.inner) + } + + fn register(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) { + datafusion::execution::memory_pool::MemoryPool::register(&self.inner, consumer) + } + + fn unregister(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) { + datafusion::execution::memory_pool::MemoryPool::unregister(&self.inner, consumer) + } + + fn grow( + &self, + reservation: &datafusion::execution::memory_pool::MemoryReservation, + additional: usize, + ) { + datafusion::execution::memory_pool::MemoryPool::grow(&self.inner, reservation, additional); + self.record(reservation.consumer().name(), additional, true); + } + + fn shrink( + &self, + reservation: &datafusion::execution::memory_pool::MemoryReservation, + shrink: usize, + ) { + datafusion::execution::memory_pool::MemoryPool::shrink(&self.inner, reservation, shrink); + self.record(reservation.consumer().name(), shrink, false); + } + + fn try_grow( + &self, + reservation: &datafusion::execution::memory_pool::MemoryReservation, + additional: usize, + ) -> std::result::Result<(), datafusion::error::DataFusionError> { + datafusion::execution::memory_pool::MemoryPool::try_grow( + &self.inner, + reservation, + additional, + )?; + self.record(reservation.consumer().name(), additional, true); + Ok(()) + } + + fn reserved(&self) -> usize { + datafusion::execution::memory_pool::MemoryPool::reserved(&self.inner) + } +} + +/// One recorded read: its path, offset, length, and the key reservation held before it. +type PeekedRead = (String, u64, u64, usize); + +/// A reader that records each read together with the key reservation held when it is issued. +struct PeekingRead { + inner: object_store::layers::mock::oio::Reader, + path: String, + reads: Arc>>, + pool: Arc, +} + +impl object_store::layers::mock::Read for PeekingRead { + async fn open( + &self, + range: object_store::layers::mock::BytesRange, + ) -> object_store::layers::mock::Result<( + object_store::layers::mock::RpRead, + Box, + )> { + let start = range.offset(); + let length = range.size().unwrap(); + self.reads.lock().unwrap().push(( + self.path.clone(), + start, + length, + self.pool.consumer("ColumnarSamplesKeys"), + )); + self.inner.open(range).await + } + + async fn read( + &self, + range: object_store::layers::mock::BytesRange, + ) -> object_store::layers::mock::Result<( + object_store::layers::mock::RpRead, + object_store::layers::mock::Buffer, + )> { + let start = range.offset(); + let length = range.size().unwrap(); + self.reads.lock().unwrap().push(( + self.path.clone(), + start, + length, + self.pool.consumer("ColumnarSamplesKeys"), + )); + self.inner.read(range).await + } +} + +/// A key the page-level window leaves over is read one at a time, and that path must charge the +/// same peak as the page-level one. +/// +/// The aligned payload cache turns one key into a whole aligned block plus whatever assembling it +/// needs. Reserving only the key entry and then reading through the cache over-commits the pool by +/// more than an order of magnitude: every key-area read must be covered by what the key +/// reservation already holds when the read is issued, whichever path serves it. +#[tokio::test] +async fn per_key_fallback_reserves_the_aligned_peak() { + use crate::cache::{CacheManager, CacheStrategy}; + use crate::series_index::format::SamplesFormatReader; + use crate::series_index::samples::reader; + use datafusion::execution::memory_pool::MemoryPool; + use object_store::layers::mock::MockLayerBuilder; + use store_api::storage::FileId; + + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + // Long keys keep the window's peak well above the pool while a single key still fits it, so + // the window is halved and the leftover keys take the per-key path. + let label = "x".repeat(1000); + let rows = 20u64; + let keys = (0..rows) + .map(|tsid| new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, tsid)) + .collect::>(); + let key_len = keys[0].len(); + let times = vec![100; keys.len()]; + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &keys, ×, false)], + ) + .await; + let (_, trailer, _, _) = root_page(&store, &handle).await; + let area = trailer.keys(); + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .page_cache_size(16 * 1024 * 1024) + .build(), + )); + let cache = strategy.columnar_payload(FileId::random()); + // What the reader charges for one retained key, plus room for one block's payload buffers. + let entry = std::mem::size_of::() + 64; + let limit = rows as usize * (key_len + entry) + 8 * 1024; + let pool = Arc::new(WatchingPool::new(limit)); + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let watching = pool.clone(); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| { + Box::new(PeekingRead { + inner, + path: path.into(), + reads: observed.clone(), + pool: watching.clone(), + }) + })) + .build() + .unwrap(); + let reader = reader::SamplesReader::open(store.clone().layer(layer), handle.clone()) + .await + .unwrap() + .with_payload_cache(cache); + let mut query = super::full_request(); + query.memory_pool = Some(pool.clone()); + reads.lock().unwrap().clear(); + let batches = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + rows as usize, + batches.iter().map(|b| b.num_rows()).sum::() + ); + // Every key-area read was paid for before it was issued: the key reservation already held at + // least the bytes the read fetches. + let key_reads = reads + .lock() + .unwrap() + .iter() + .filter(|(path, start, length, _)| { + path.ends_with(".samples") && *start < area.end && *start + *length > area.start + }) + .cloned() + .collect::>(); + assert!( + key_reads.len() > 1, + "{} key-area reads means the leftover keys never took the per-key path", + key_reads.len() + ); + for (_, start, length, reserved) in &key_reads { + assert!( + *reserved >= *length as usize, + "a {length} byte key-area read at {start} was issued while the key reservation held \ + only {reserved} bytes" + ); + } + // The pool has to hold every retained key, and the scan never asks for more than it has. + assert!( + pool.high_water() >= rows as usize * (key_len + entry), + "high water {} is below the {} bytes of retained keys", + pool.high_water(), + rows as usize * (key_len + entry) + ); + assert_eq!(0, pool.reserved()); + + // A pool that cannot pay for one key's peak must fail before any key-area read, even with the + // aligned cache enabled. + let pool = Arc::new(WatchingPool::new(key_len)); + let mut query = super::full_request(); + query.memory_pool = Some(pool.clone()); + reads.lock().unwrap().clear(); + let error = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .expect_err("a pool below one key's peak must refuse the read"); + assert_eq!( + common_error::status_code::StatusCode::RuntimeResourcesExhausted, + common_error::ext::ErrorExt::status_code(&error) + ); + let key_reads = reads + .lock() + .unwrap() + .iter() + .filter(|(path, start, length, _)| { + path.ends_with(".samples") && *start < area.end && *start + *length > area.start + }) + .count(); + assert_eq!( + 0, key_reads, + "{key_reads} key-area reads while the pool cannot pay" + ); + assert_eq!(0, pool.reserved()); +} + +/// A series that spans two parts is recognised by the last loaded key of the first part. +/// +/// That copy outlives every per-part reservation, so it has a consumer of its own: charged while +/// the boundary key is alive and given back once the boundary is behind the scan. Copying it from +/// a block's key would retain bytes no reservation pays for. +#[tokio::test] +async fn boundary_key_is_charged_while_it_is_alive() { + use crate::series_index::format::SamplesFormatReader; + use crate::series_index::samples::reader; + use datafusion::execution::memory_pool::MemoryPool; + + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat { + block_rows: 1, + part_bytes: 1, + ..Default::default() + }; + let label = "x".repeat(1024); + let key = new_sparse_primary_key(&[label.as_str(), "z"], &meta, 1, 7); + // One series, two blocks, one block per part: the series continues across the boundary. + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch(&meta, &vec![key.clone(); 2], &[10, 20], false)], + ) + .await; + assert_eq!(2, super::manifest(&store, &handle).await.files.len()); + let pool = Arc::new(WatchingPool::new(8 * 1024 * 1024)); + let reader = reader::SamplesReader::open(store, handle).await.unwrap(); + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.memory_pool = Some(pool.clone()); + let mut stream = reader.read(query).unwrap(); + // The first part is not finished yet: nothing has been carried across a boundary. + assert!( + stream.try_next().await.unwrap().is_some(), + "the first part holds one block" + ); + assert_eq!( + 0, + pool.consumer("ColumnarSamplesBoundaryKey"), + "the boundary key must not be charged before the part ends" + ); + // The second batch is only produced after the first part handed its last loaded key over. + assert!( + stream.try_next().await.unwrap().is_some(), + "the series continues in the second part" + ); + let charged = pool.consumer("ColumnarSamplesBoundaryKey"); + assert!( + charged >= key.len(), + "the boundary key's {charged} charged bytes do not cover its {} bytes", + key.len() + ); + // Past the last part the boundary key is no longer needed, so its charge goes back. + assert!(stream.try_next().await.unwrap().is_none()); + assert_eq!(0, pool.consumer("ColumnarSamplesBoundaryKey")); + drop(stream); + assert_eq!(0, pool.reserved()); +} + +/// A selection that prunes a block must not cut the order check between the keys around it. +/// +/// The pruned block's key is never read, so the check has to chain over the last *loaded* key. +/// Comparing against the last visited block instead silently drops the comparison, and two +/// reversed exact keys under one `(table_id, tsid)` pass as ordered. +#[tokio::test] +async fn skipped_block_does_not_hide_reversed_keys() { + let store = store(); + let meta = super::metadata(); + let format = ColumnarSamplesFormat::default(); + let first = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let middle = new_sparse_primary_key(&["host-000002", "z"], &meta, 1, 7); + let last = new_sparse_primary_key(&["host-000003", "z"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[first, middle, last], + &[100, 5_000, 110], + false, + )], + ) + .await; + // Swapping the first and last series' stored entries keeps both checksums and identifiers + // valid; only the order of the two keys the request loads is wrong. The middle block is + // outside the requested time range, so its key is never read. + let (manifest, _, page, mut bytes) = root_page(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the fixture") + }; + assert_eq!(3, blocks.len()); + let entries = [&blocks[0].key, &blocks[2].key] + .map(|reference| reference.offset as usize..reference.checksum_end().unwrap() as usize); + let first_entry = bytes[entries[0].clone()].to_vec(); + bytes.copy_within(entries[1].clone(), entries[0].start); + bytes[entries[1].clone()].copy_from_slice(&first_entry); + store.write(&manifest.part_path(0), bytes).await.unwrap(); + let reader = format.reader(store, handle).await.unwrap(); + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.end = Some(1_000); + let error = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .expect_err("a skipped block must not hide reversed keys"); + assert!( + error + .to_string() + .contains("unordered sample exact series keys"), + "unexpected error: {error}" + ); +} + +/// The same gap across a part boundary: a part whose last block was pruned must still hand its +/// last *loaded* key to the next part. +/// +/// The skipped block ends the first part, so a boundary key taken from "the last visited block" +/// is missing and the check would silently compare nothing across the parts. +#[tokio::test] +async fn pruned_block_does_not_hide_the_boundary_key() { + let store = store(); + let meta = super::metadata(); + // One leaf holds two descriptors and the root holds one leaf, so the first two blocks share a + // page in the first part and the third block starts the second part. The pruned block is + // therefore traversed without being loaded, which is what used to lose the boundary key. + let format = ColumnarSamplesFormat { + block_rows: 1, + page_bytes: 476, + root_capacity: 1, + ..Default::default() + }; + let first = new_sparse_primary_key(&["host-000001", "z"], &meta, 1, 7); + let middle = new_sparse_primary_key(&["host-000002", "z"], &meta, 1, 7); + let last = new_sparse_primary_key(&["host-000003", "z"], &meta, 1, 7); + let handle = build( + &format, + &store, + meta.clone(), + &[keyed_batch( + &meta, + &[first, middle, last.clone()], + &[100, 5_000, 110], + false, + )], + ) + .await; + assert_eq!(2, super::manifest(&store, &handle).await.files.len()); + // Rewrite the second part's key with a smaller key of the same length and identifier, so the + // two loaded keys are reversed while everything else about the part stays valid. + let smaller = new_sparse_primary_key(&["host-000000", "z"], &meta, 1, 7); + assert_eq!(last.len(), smaller.len()); + let (_, page, mut bytes) = part_root(&store, &handle, 1).await; + let manifest = super::manifest(&store, &handle).await; + let Page::Leaf(blocks) = page else { + panic!("one leaf page holds the second part") + }; + let reference = blocks[0].key; + let start = reference.offset as usize; + bytes[start..start + smaller.len()].copy_from_slice(&smaller); + let sum = codec::checksum(&smaller); + bytes[start + smaller.len()..reference.checksum_end().unwrap() as usize] + .copy_from_slice(&sum.to_le_bytes()); + store.write(&manifest.part_path(1), bytes).await.unwrap(); + let reader = format.reader(store, handle).await.unwrap(); + let mut query = request(&[TABLE_ID_COLUMN, TSID_COLUMN, "ts"]); + query.end = Some(1_000); + let error = reader + .read(query) + .unwrap() + .try_collect::>() + .await + .expect_err("a pruned block must not hide the boundary key"); + assert!( + error + .to_string() + .contains("unordered sample exact series keys"), + "unexpected error: {error}" + ); +} diff --git a/src/mito2/src/series_index/samples/writer.rs b/src/mito2/src/series_index/samples/writer.rs index e2d9fa174aa..7ce8d9e5588 100644 --- a/src/mito2/src/series_index/samples/writer.rs +++ b/src/mito2/src/series_index/samples/writer.rs @@ -32,15 +32,16 @@ use crate::error::{Result, SerdeJsonSnafu}; use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata, ArtifactObject}; use crate::series_index::artifact_writer::Outputs; use crate::series_index::format::FormatWriter; +use crate::series_index::identity::{ExactSeriesKey, IdentityEncoding, decode_series_id}; use crate::series_index::samples::codec; use crate::series_index::samples::metadata::{ - self, Block, Column, Page, PageReference, Trailer, pack, + self, Block, Column, KeyRef, Page, PageReference, Trailer, pack, }; use crate::series_index::samples::{ ColumnarSamplesFormat, MANIFEST_LIMIT, Manifest, MetricSeriesId, Part, TimestampColumn, TimestampUnit, VERSION, ValueColumn, invalid, require, valid_rows, }; -use crate::series_index::writer::{decode_series_identity, validate_metadata}; +use crate::series_index::writer::validate_metadata; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; @@ -82,22 +83,101 @@ struct Row { operation: u8, fields: Vec>, } +/// A block descriptor whose exact key has not been placed in the part's key area yet. +/// +/// The key bytes are held here until the part ends, when they are written to their own area +/// and replaced by a fixed-size [`KeyRef`]. That is why no descriptor — and therefore no +/// metadata page — ever has to fit a key. +struct PendingBlock { + key: ExactSeriesKey, + min: i64, + max: i64, + ordinal: u64, + rows: u32, + columns: Vec, +} + +impl PendingBlock { + /// Descriptor bytes once the key reference replaces the key bytes. + fn size(&self) -> usize { + metadata::descriptor_bytes(self.columns.len()) + } +} + +/// Places one copy of each series key in the part's key area and turns pending blocks into +/// descriptors that reference it. +/// +/// Blocks are ordered by exact key, so every block of one series is adjacent and shares one +/// copy of the key. Consecutive equal keys therefore reuse the previous reference, while every +/// new series appends its bytes after the previous copy: references are ordered, disjoint, and +/// unique per series, which is what lets a reader treat an equal reference as the same series +/// without reading the bytes. Each copy is followed by its XXH64, which the reference does not +/// count: a reader verifies the bytes before they can name a series. +fn place_keys(leaves: Vec>, start: u64) -> Result<(Vec>, Vec)> { + let mut area: Vec = Vec::new(); + let mut previous: Option<(Vec, KeyRef)> = None; + let mut placed = Vec::with_capacity(leaves.len()); + for pending in leaves { + let mut blocks = Vec::with_capacity(pending.len()); + for block in pending { + let reuse = previous + .as_ref() + .filter(|(key, _)| key.as_slice() == block.key.bytes()) + .map(|(_, reference)| *reference); + let reference = match reuse { + Some(reference) => reference, + None => { + let length = u32::try_from(block.key.bytes().len()) + .map_err(|_| invalid("sample series key exceeds 4 GiB"))?; + let offset = start + .checked_add( + u64::try_from(area.len()) + .map_err(|_| invalid("sample key area overflow"))?, + ) + .ok_or_else(|| invalid("sample key offset overflow"))?; + area.extend_from_slice(block.key.bytes()); + area.extend_from_slice(&codec::checksum(block.key.bytes()).to_le_bytes()); + let reference = KeyRef { offset, length }; + previous = Some((block.key.bytes().to_vec(), reference)); + reference + } + }; + blocks.push(Block { + id: block.key.id(), + key: reference, + min: block.min, + max: block.max, + ordinal: block.ordinal, + rows: block.rows, + columns: block.columns, + }); + } + placed.push(blocks); + } + Ok((placed, area)) +} pub(super) struct SamplesWriter { options: ColumnarSamplesFormat, manifest: Manifest, outputs: Outputs, output_buffer: SampleBuffer, offsets: SparseOffsetsCache, - key_buffer: Vec, + /// Canonical sparse key of the series that owns `rows` and `last`. + series_key: Option, rows: Vec, last: Option<(MetricSeriesId, i64, u64)>, ordinal: u64, - leaves: Vec>, + leaves: Vec>, leaf_bytes: usize, descriptor_bytes: usize, position: u64, part_rows: u64, part_blocks: u64, + /// Bytes this part's key area is charged, one copy per series with its checksum, counted + /// before they are written. + part_key_bytes: u64, + /// The last series whose key this part already counts, to charge one copy per series. + part_last_key: Option, manifest_bytes: usize, failed: bool, } @@ -142,6 +222,7 @@ impl SamplesWriter { let manifest = Manifest { format: ArtifactFormat::ColumnarSamples, version: VERSION, + identity_encoding: IdentityEncoding::CURRENT_SAMPLES, region_id: metadata.region_id.as_u64(), source_file_id, generation, @@ -174,7 +255,7 @@ impl SamplesWriter { }, output_buffer: SampleBuffer::default(), offsets: SparseOffsetsCache::new(), - key_buffer: Vec::new(), + series_key: None, rows: Vec::new(), last: None, ordinal: 0, @@ -184,6 +265,8 @@ impl SamplesWriter { position: 0, part_rows: 0, part_blocks: 0, + part_key_bytes: 0, + part_last_key: None, manifest_bytes, failed: false, }) @@ -248,27 +331,49 @@ impl SamplesWriter { )?; values.value(index) }; - // TODO: coordinate exact keys and future job/instance identity components - // with the series format. This port assumes (table_id, tsid) is unique. - let (id, _) = - decode_series_identity(key, &[], &mut self.offsets, &mut self.key_buffer)?; let time = times.value(i); let seq = sequence.value(i); let op = operations.value(i); require(op <= 1, "unknown sample operation")?; - if let Some((last_id, last_time, last_seq)) = self.last { - require( - id >= last_id - && (id != last_id - || time > last_time - || (time == last_time && seq <= last_seq)), - "sample input is not sorted", - )?; - if id != last_id { + // Identity is the exact sparse key, not `(table_id, tsid)`: two distinct + // series may share a `tsid`, and every block records which exact key owns it. + let new_series = self + .series_key + .as_ref() + .is_none_or(|current| current.bytes() != key); + if new_series { + // Validate the whole key once per series: identity is the exact key, and a + // malformed one must not reach a block descriptor. + let exact = ExactSeriesKey::from_sparse_key(key, &mut self.offsets)?; + if self.last.is_some() { + // Rows arrive in canonical sparse-key order. + let previous = self + .series_key + .as_ref() + .ok_or_else(|| invalid("missing sample series key"))?; + require( + exact > *previous, + "sample input is not sorted by primary key", + )?; + // A different key that shares `(table_id, tsid)` is a different + // series: it opens its own block, which records its exact key, so + // timestamps and ordinals may restart. self.flush_block().await?; self.ordinal = 0; } + self.series_key = Some(exact); } + let id = match self.last { + Some((id, last_time, last_seq)) if !new_series => { + // The same exact series continues: the existing monotonic rules hold. + require( + time > last_time || (time == last_time && seq <= last_seq), + "sample input is not sorted", + )?; + id + } + _ => decode_series_id(key, &mut self.offsets)?, + }; if self.rows.len() == self.options.block_rows || !valid_rows(self.rows.len() + 1, fields.len()) { @@ -299,12 +404,22 @@ impl SamplesWriter { return Ok(()); } let rows = std::mem::take(&mut self.rows); - let (mut block, chunks) = encode_block(&rows, &self.manifest, self.ordinal)?; + let key = self + .series_key + .as_ref() + .ok_or_else(|| invalid("sample block has no series key"))?; + let (mut block, chunks) = encode_block(&rows, key, &self.manifest, self.ordinal)?; + // Bytes one copy of this series' key costs in a part's key area: the key and its stored + // checksum. + let key_len = key.bytes().len() as u64 + metadata::KEY_CHECKSUM_BYTES; self.ordinal = self .ordinal .checked_add(1) .ok_or_else(|| invalid("sample ordinal overflow"))?; let size = block.size(); + // The descriptor carries a fixed-size key reference, not the key, so this bound is a + // function of the field count only: a key longer than a whole metadata page is no + // longer a build failure, and only a narrower configured page can reject a descriptor. require( size + 12 <= self.options.page_bytes, "sample descriptor exceeds leaf limit", @@ -344,7 +459,18 @@ impl SamplesWriter { .last_mut() .ok_or_else(|| invalid("missing sample leaf"))? .push(block); - if self.position >= self.options.part_bytes + // Keys live in their own area but are still bytes of this part: count one copy per + // series so the part target covers them too, and a key-heavy part still rolls over. + // This is decided after the rollover above, which cleared the part's accounting: a series + // that starts a new part still needs its copy counted in that part. + if self.part_last_key.as_ref() != self.series_key.as_ref() { + self.part_key_bytes = self + .part_key_bytes + .checked_add(key_len) + .ok_or_else(|| invalid("sample key area overflow"))?; + self.part_last_key = self.series_key.clone(); + } + if self.position.saturating_add(self.part_key_bytes) >= self.options.part_bytes || self.descriptor_bytes >= self.options.metadata_bytes { self.finish_part().await?; @@ -363,8 +489,24 @@ impl SamplesWriter { return Ok(()); } let payload_end = self.position; + // Keys go to their own area after the payload. Every block of one series shares one + // copy, so the area holds one key per distinct series of this part. + let (leaves, key_area) = place_keys(std::mem::take(&mut self.leaves), payload_end)?; + if !key_area.is_empty() { + self.position = self + .position + .checked_add( + u64::try_from(key_area.len()) + .map_err(|_| invalid("sample key area overflow"))?, + ) + .ok_or_else(|| invalid("sample file length overflow"))?; + self.output_buffer + .write(&mut self.outputs, &key_area) + .await?; + } + let key_end = self.position; let mut children = Vec::new(); - for leaf in std::mem::take(&mut self.leaves) { + for leaf in leaves { children.push(self.write_page(Page::Leaf(leaf)).await?); } let depth = if children.len() == 1 { 1 } else { 2 }; @@ -386,6 +528,7 @@ impl SamplesWriter { depth, file_length, payload_end, + key_end, } .encode(), ) @@ -413,6 +556,8 @@ impl SamplesWriter { self.position = 0; self.part_rows = 0; self.part_blocks = 0; + self.part_key_bytes = 0; + self.part_last_key = None; self.descriptor_bytes = 0; self.leaf_bytes = 12; Ok(()) @@ -452,6 +597,23 @@ impl SamplesWriter { }) } } + +#[cfg(test)] +impl SamplesWriter { + /// The current part's accounting, for the rollover tests: `(position, counted key bytes, + /// blocks in the current part, finished parts)`. + /// + /// A key copy that is not counted defers the part boundary silently, and the parts it makes + /// are still readable, so only the accounting itself shows the difference. + pub(super) fn part_accounting(&self) -> (u64, u64, u64, usize) { + ( + self.position, + self.part_key_bytes, + self.part_blocks, + self.manifest.files.len(), + ) + } +} #[async_trait] impl FormatWriter for SamplesWriter { async fn write(&mut self, batch: &RecordBatch) -> Result<()> { @@ -476,11 +638,20 @@ impl FormatWriter for SamplesWriter { self.outputs.cleanup().await } } -fn encode_block(rows: &[Row], manifest: &Manifest, ordinal: u64) -> Result<(Block, Vec>)> { +fn encode_block( + rows: &[Row], + key: &ExactSeriesKey, + manifest: &Manifest, + ordinal: u64, +) -> Result<(PendingBlock, Vec>)> { let first = rows.first().ok_or_else(|| invalid("empty sample block"))?; let last = rows.last().ok_or_else(|| invalid("empty sample block"))?; - let mut block = Block { - id: first.id, + require( + rows.iter().all(|row| row.id == key.id()), + "sample block mixes series", + )?; + let mut block = PendingBlock { + key: key.clone(), min: first.time, max: last.time, ordinal, diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index 28999d8409b..21731505238 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -27,7 +27,6 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::timestamp::timestamp_array_to_primitive; use mito_codec::row_converter::SparseOffsetsCache; -use mito_codec::row_converter::sparse::SparsePrimaryKeyView; use object_store::ObjectStore; use parquet::file::metadata::KeyValue; use snafu::{OptionExt, ResultExt, ensure}; @@ -36,9 +35,7 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; use store_api::storage::consts::ReservedColumnId; -use crate::error::{ - DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, -}; +use crate::error::{InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result}; use crate::series_index::{ MAX_TS_COLUMN, MIN_TS_COLUMN, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, }; @@ -595,9 +592,9 @@ fn decode_primary_key( offsets: &mut SparseOffsetsCache, buf: &mut Vec, ) -> Result { - let (id, tags) = decode_series_identity(primary_key, tag_columns, offsets, buf)?; - let table_id = id.table_id; - let tsid = id.tsid; + let (key, tags) = decode_series_identity(primary_key, tag_columns, offsets, buf)?; + let table_id = key.id().table_id; + let tsid = key.id().tsid; Ok(SeriesIndexRow { min_ts, @@ -610,27 +607,10 @@ fn decode_primary_key( } /// Shared sparse-key extraction, independent of timestamp aggregation. -pub(crate) fn decode_series_identity( - primary_key: &[u8], - tag_columns: &[(ColumnId, String)], - offsets: &mut SparseOffsetsCache, - buf: &mut Vec, -) -> Result<(crate::series_index::MetricSeriesId, Vec>)> { - let mut view = SparsePrimaryKeyView::new(primary_key, offsets).context(DecodeSnafu)?; - let id = crate::series_index::MetricSeriesId { - table_id: view.table_id(), - tsid: view.tsid(), - }; - let mut tags = Vec::with_capacity(tag_columns.len()); - for (column_id, _) in tag_columns { - tags.push( - view.label(*column_id, buf) - .context(DecodeSnafu)? - .map(str::to_owned), - ); - } - Ok((id, tags)) -} +/// +/// The exact key it returns is what decides series identity; `MetricSeriesId` is +/// only its coarse, hash-based prefix. +pub(crate) use crate::series_index::identity::decode_series_identity; /// Builds a timestamp array in `unit` from raw i64 values. fn ts_array(timestamps: impl Iterator, unit: TimeUnit) -> ArrayRef { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index ce756c81828..5b29a7bf373 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -531,8 +531,14 @@ async fn series_index_store_from_config( Ok(Some(store)) } -/// Bump when the rebuildable local index layout becomes incompatible. -const SERIES_INDEX_CACHE_VERSION: &str = "2"; +/// Marks the rebuildable local index layout, and resets the store when it changes. +/// +/// This branch is experimental and has no deployments, so it pins the token at `1`: a +/// cache written by any other build is discarded on startup and rebuilt from the SSTs. +/// Artifacts are additionally rejected on open when they lack the mandatory +/// `identity_encoding` discriminator, so a stale artifact is never served even if the +/// reset did not happen (a read that precedes it, or an artifact outside this cache). +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<()> { diff --git a/tests/compatibility/wire/mito_columnar_samples/README.md b/tests/compatibility/wire/mito_columnar_samples/README.md index 490aa343a69..726fd076e9b 100644 --- a/tests/compatibility/wire/mito_columnar_samples/README.md +++ b/tests/compatibility/wire/mito_columnar_samples/README.md @@ -5,7 +5,7 @@ This fixture freezes the version-1 sample wire contract in the owning Mito Rust test, not SQLness discovery: ```shell -cargo nextest run -p mito2 frozen_v1_samples_fixture_and_unknown_codec +cargo nextest run -p mito2 frozen_v1_samples_fixture_is_rejected ``` `v1-manifest.json` and `v1-part.hex` describe a 397-byte part with raw compression, @@ -26,11 +26,23 @@ little-endian fields and XXH64 with seed zero. Its layout is: | 53..285 | Leaf/root: one block, four column references | | 285..397 | Trailer with root reference and checksum | -The test installs these exact bytes and recorded object sizes, opens them through -`SamplesFormat`, and verifies the values and order. It also changes a column -codec to an unknown ID while updating checksums to test that unrequested columns -can be skipped and requested unknown codecs fail explicitly. Legacy catalog -identifier coverage remains in `series_index::catalog` tests. +The current reader writes the same wire version as this fixture, so the version +number alone does not reject it. The manifest carries no `identity_encoding` +discriminator, and the part records only `(table_id, tsid)`, which is a hash: a +part whose two series were merged is structurally valid and ordered exactly like a +correct one. The test installs these exact bytes and recorded object sizes and +asserts that opening them through `SamplesFormat` fails on the missing +`identity_encoding` discriminator, which is what makes maintenance rebuild the +artifact from its SSTs instead of reusing it. Legacy catalog identifier coverage +remains in `series_index::catalog` tests. + +The part bytes above also predate two later version-1 shapes, neither of which is +migrated: block descriptors now carry a fixed-size key reference into a separate key +area instead of the key bytes, and the trailer records that area (`key_end`) and is +120 bytes, not 112. The fixture is still rejected for the same reason — the +discriminator is missing — and that rejection happens while opening the manifest, +before any part byte is decoded, so the frozen bytes are not a reader for the current +part layout and must not be regenerated to look like one. Do not regenerate this fixture just to make a reader change pass. Preserve its bytes for future reader compatibility; introduce a separate versioned fixture