diff --git a/docs/formats/mito-samples.md b/docs/formats/mito-samples.md new file mode 100644 index 00000000000..304977a7893 --- /dev/null +++ b/docs/formats/mito-samples.md @@ -0,0 +1,650 @@ +# Mito columnar samples format + +**Status: implemented as an opt-in internal backend.** +This document specifies a standalone samples writer and reader. Query-consumer +integration, caches, and enabling samples by default are separate work. + +The port is stacked on [#9841](https://github.com/GreptimeTeam/greptimedb/pull/9841), +at commit `6b44040c0ad96fb06d592c7bba5a1744e2274718`, and follows the presentation +and artifact conventions in [Mito columnar series](mito-series.md). Its sample +encodings and single-series blocks derive from the benchmark's +[columnar samples v2 specification][bench-spec] and [implementation][bench-code] +at revision `8650c3003a861c8cf2c9f5037bfc81fde6a3bdc4`. +The benchmark variant name does not determine the Mito wire version. + +The principal format choices are nullable Float64 fields, preservation of source +rows and operation types, blocks bounded by rows and bytes without hourly splits, +and Mito framing with sample-specific time bounds. This is an independent format; +it does not read or write benchmark files. + +## Artifact manifest + +The catalog identifier is `columnar-samples`. Its file and manifest version is 1, +independent of the series format's version and both formats' codec identifiers. +Existing `parquet` and `columnar-series` identifiers keep their meaning. + +One artifact is derived from one immutable source SST. It owns these objects: + +```text +/samples///manifest.json +/samples///part-0.samples +/samples///part-1.samples +... +``` + +The source file ID and generation use existing `FileId` serialization. The +generation is unique for each build. A manifest contains: + +| Field | Meaning | +| --- | --- | +| `format` | String `columnar-samples` | +| `version` | Integer 1 | +| `region_id` | Physical region ID, u64 | +| `source_file_id` | SST `FileId` matching `SamplesEntry.file_id` | +| `generation` | `FileId` matching the catalog artifact generation | +| `timestamp` | Object with `column_id`, `name`, and `unit` | +| `fields` | Ordered objects with `column_id`, `name`, `type`, and `nullable` | +| `files` | Ordered part descriptors, defined below | + +The timestamp's ID is u32, its name is the region's time-index name, and its unit +is one of `second`, `millisecond`, `microsecond`, or `nanosecond`. Values are +signed i64 counts since the Unix epoch in that unit. There is no dataset-relative +origin, unit conversion, or timezone in the file. The timestamp is non-null. + +The field list records the build-time region field schema in metadata order. +Require 1..=1,024 fields, each with `type = "float64"` and `nullable = true`. +A non-nullable source Float64 field is represented by this same nullable logical +type with no null values. Sparse logical-table fields may be entirely null. +Reject any other source field type; do not silently drop or cast it. Timestamp +and field IDs must be distinct, and all output names must be unique, including +the internal names below. Tags are not sample fields and are not stored here. + +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 exposed by the format reader. +Schema reconciliation with a later region schema belongs to a future consumer; +the reader always exposes the stored schema. + +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`. +`first` and `last` are inclusive `{ table_id: u32, tsid: u64 }` bounds. Time bounds +are the minimum and maximum over **all** blocks in the part, not just its first +and last rows. Every part is nonempty. Part IDs start at zero, increase by one, +and derive `part-.samples` names with decimal IDs and no leading zeros. +No object path is stored in a part descriptor. + +Parts follow input order. Adjacent identifier ranges satisfy +`previous.last <= next.first`; equality is allowed because a series can continue +in another part. The shared boundary series can also have equal timestamps on +both sides. These ranges are **not disjoint**, unlike the series-index format. +Readers consider every matching part rather than choosing a unique owner. + +The manifest is written only after all parts close successfully. Existing samples +catalog publication commits the artifact. `ArtifactMetadata.objects` records the +manifest and every part as `{ path, file_size }`; the artifact's `file_size` is +their checked sum. Record sizes while writing, including the manifest's size. +An empty writer produces a manifest with an empty file list and no part files; +the existing builder may instead skip an SST for which it has no readable input. + +On open, bound the manifest to 64 MiB before parsing. Validate its format, schema, +part order, and sizes; match region, source SST, and generation to the pinned handle. +Require the artifact object list to contain exactly the derived manifest and part +paths, without duplicates, extra objects, or inconsistent size totals. Read order +comes from the manifest, not the lifecycle object list. + +## Part file + +```text +┌──────────────────────────────────────────────┐ +│ Header 16 bytes │ +├──────────────────────────────────────────────┤ +│ Independently encoded/compressed chunks │ +│ Block 0: time, sequence, operation, fields │ +│ Block 1: ... │ +├──────────────────────────────────────────────┤ ← payload_end +│ Leaf pages (one is the root at depth 1) │ +│ Optional internal root (depth 2) │ +├──────────────────────────────────────────────┤ +│ Trailer 112 bytes │ +└──────────────────────────────────────────────┘ +``` + +All fixed-width numbers are little-endian, without alignment padding. Byte ranges +are half-open; offsets are absolute from the start of the part. Parse each field +explicitly. Rust struct layout, enum order, and `usize` are not wire contracts. +Reserved fields must be zero, and decoded structures must have no trailing bytes. + +### Header + +| Field | Bytes | Value | +| --- | ---: | --- | +| Magic | 8 | ASCII `GTSAMPLE` | +| File version | 4 | 1 | +| Reserved | 4 | Zero | + +### 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 +are independent chunks; their references are in the metadata, not inline headers. +Emit chunks in that order, with value fields following manifest order. + +The input is the flat `RecordBatch` stream from the existing per-SST samples +builder, after its `FlatCompatBatch` schema adaptation. Preserve every row +delivered by that builder. The writer neither deduplicates timestamps nor chooses +the highest sequence. It preserves delete markers, duplicate rows, and equal-key +ties, including conflicting values at the same sequence. This is not a statement +that an SST itself retains every historical write. + +Require nondecreasing identifiers, nondecreasing timestamps within an identifier, +and nonincreasing sequence within equal timestamps, across batches and parts. +Equal ordering keys retain input order; operation type is not a sorting key. +Reject decreasing input rather than sorting or buffering a whole SST. Future +query adapters remain responsible for snapshot visibility, cross-SST merging, +deduplication, and delete processing. + +`(table_id, tsid)` is the complete identity under #9841's current uniqueness +assumption. Equal TSIDs in different tables are distinct. Exact sparse primary +keys and tags are not retained, and this format does not promise collision-safe +identity. TODO: evaluate exact-key storage and collision handling with the series +format before using these sidecars as a general query source. + +TODO: evaluate adding dedicated `job_id` and `instance_id` components to the +identifier for locality, following the VM native approach. The current +`(table_id, tsid)` pair is the version-1 identity, not a permanent restriction on +future formats. The extended identifier's component order, widths, derivation, +and relationship to TSID uniqueness require a separate design and benchmarks. +Coordinate that change with the series format: sample block identities, +part/page bounds, sorting, and selected-series requests must agree. Changing +the persisted identifier layout requires a new format version; do not repurpose +reserved bytes or silently reinterpret existing keys. Ordinary `job` and +`instance` tags do not become dedicated identifier components in this port. + +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 +limit is 8,192, also the version-1 reader maximum; a writer may choose a smaller +positive limit without changing the version. With `F` manifest fields and `N` +rows, the decoded accounting size is: + +```text +N * (8 timestamp + 8 sequence + 1 operation + 8 * F values) + + F * ceil(N / 8) validity bytes +``` + +Require this to be at most 262,144 bytes. Count all manifest fields, including +omitted/all-null fields, and reserve validity bytes even for all-valid fields. +This is a bound on decoded column buffers, not total allocator usage or RSS. +The effective row cap is the largest integer satisfying both limits: 8,192 for +one field and 31 for 1,024 fields. Encoded buffers and metadata have separate +bounds below. Increasing the reader maximum requires an explicit compatibility +decision; it is not merely a writer option in this version. + +Materializing the two identity columns adds at most `12 * N` bytes to an output +batch. Thus a full-schema output batch has at most `262144 + 12 * 8192 = 360448` +bytes of value/validity buffers under this accounting, plus Arrow and allocation +overhead. Temporary encoded, compact non-null, and selected-output buffers can +coexist; the block limit is not a promise that total reader memory is 256 KiB. + +### Payload reference + +Every stored column and metadata page uses the same 32-byte reference layout as +the series format: + +| Field | Bytes | Meaning | +| --- | ---: | --- | +| Offset | 8 | Absolute location of stored bytes | +| Stored length | 4 | Bytes to fetch | +| Decoded length | 4 | Bytes after decompression, before column/page decoding | +| Checksum | 8 | XXH64 of stored bytes, seed zero | +| Encoding ID | 2 | Sample column or metadata encoding, according to context | +| Compression ID | 2 | Compression of the encoded bytes | +| Reserved | 4 | Zero | + +Lengths are nonzero. Check `offset + stored_length` for u64 overflow before +checking the enclosing area. u32 payload lengths do not limit the part or its +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. + +### Metadata pages + +A leaf describes blocks. An internal page references leaves. The root is either +the only leaf (depth 1) or an internal page (depth 2). Other depths are unsupported +in version 1; roll parts before a third level would be needed. + +As in the series format, the entire page, including its kind and count, is encoded +and independently compressed. The reference in the parent or trailer supplies +its storage location, lengths, checksum, and codecs. There is no separate stored +page header. Decoded pages begin with `kind: u32` and `entry_count: u64`: + +```text +Leaf (kind 0): Internal root (kind 1): + kind u32 kind u32 + entry_count u64 entry_count u64 + block descriptor 0 page reference 0 72 bytes + ... ... + block descriptor N-1 page reference N-1 72 bytes +``` + +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. + +#### Block descriptor + +| Field | Bytes | Meaning | +| --- | ---: | --- | +| Table ID | 4 | Constant for this block | +| TSID | 8 | Constant for this block | +| 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 | +| Row count | 4 | Nonzero, subject to decoded-block limits | +| 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. + +Each column descriptor is: + +| Field | Bytes | Meaning | +| --- | ---: | --- | +| Role | 4 | 1 timestamp, 2 sequence, 3 operation type, 4 value field | +| Column ID | 4 | Timestamp/field physical ID; zero for sequence/operation | +| Payload reference | 32 | Independently stored chunk | +| Null count | 4 | Number of null rows | + +Role and column ID together identify a chunk. The role namespace avoids reserving +physical field IDs for internal columns. Time, sequence, and operation descriptors +are mandatory, appear first in that order, and have zero null counts. Value +descriptors follow manifest order with no repeated or unknown IDs. A value with +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 in this port. + +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. + +#### Page reference and time bounds + +Samples extend the series format's reference with two timestamp bounds: + +| Field | Bytes | Meaning | +| --- | ---: | --- | +| First table ID; first TSID | 4 + 8 | Inclusive first identifier | +| Last table ID; last TSID | 4 + 8 | Inclusive last identifier | +| Minimum timestamp | 8 | Minimum over all descendant blocks | +| Maximum timestamp | 8 | Maximum over all descendant blocks | +| Payload reference | 32 | Location and codecs of the child page | + +The reference is 72 bytes. Compare identifiers as integer pairs, not as serialized +little-endian bytes. Time bounds use the manifest unit. The time range is an +envelope over all identifiers in the page; it need not describe the first or last +identifier alone. This can admit extra reads, but cannot exclude matching rows. +Every part/page identifier interval requires `first <= last`, and all part/page/ +block time intervals require `min_timestamp <= max_timestamp`, even when pruning +would skip their payloads. + +Child pages follow block order. Their identifier bounds may touch at a shared +series, just like part bounds. Do not require strict disjointness or use an +exclusive next-child lower bound. Validate each loaded page's derived identifier +and time bounds against its reference. At depth 2 all children must be leaves; +reject an internal child. Child byte ranges are ordered, disjoint, within the +metadata area, and precede the parent. The root ends immediately before the trailer. + +An internal root has `12 + 72 * N` decoded bytes, so it holds at most +`floor((65536 - 12) / 72) = 910` leaf references. The writer must account for the +next descriptor's leaf placement before writing its block. If the descriptor +would require a 911th leaf, finish this part and start the block in a new part. + +#### Navigation and equal-timestamp example + +For a requested qualified identity and optional half-open interval `[start, end)`, +select every part/page whose identifier interval contains the identity and whose +time envelope overlaps the request: + +```text +(start is absent || max_timestamp >= start) + && (end is absent || min_timestamp < end) +``` + +For multiple selected identities, intersect the ordered ID set with the inclusive +identifier range. A table restriction also intersects this selection. Visit each +selected page once per request. Apply the same overlap test to leaf blocks, then +filter decoded timestamps exactly. Metadata bounds are conservative candidates. +No predecessor-only or unique-child lookup is used: a shared identity or timestamp +can require several adjacent children, and even several parts. + +```text +Part 0, depth 2: + Root R + ├─ Leaf A: IDs [(7,42),(7,42)], times [0,100] + │ block ordinal 0: (7,42), times [0,100] + └─ Leaf B: IDs [(7,42),(7,43)], times [0,200] + block ordinal 1: (7,42), times [100,200] + block ordinal 0: (7,43), times [0,50] + +Part 1: + Leaf/root C: IDs [(7,43),(7,43)], times [50,80] + block ordinal 1: (7,43), times [50,80] +``` + +For `(7,42)` and `[100,101)`, visit both A and B, keeping all rows at timestamp +100 from blocks 0 and 1. Do not decode `(7,43)` payloads. For `(7,43)` and +`[50,51)`, visit B and part 1's C: both may contain rows at 50. A `[51,60)` request +for `(7,43)` can skip B's `(7,43)` block even though B's time envelope overlaps. +For `(7,42)` and `[90,95)`, A is still visited although its block starts before +the requested interval. Global time envelopes must not be computed from just +the first and last blocks: B's range is `[0,200]`, not `[100,50]`. + +This differs from #9841, which has no sample timestamps or time navigation. + +### Trailer + +| Field | Bytes | Meaning | +| --- | ---: | --- | +| Root page reference | 72 | Inclusive ID/time bounds and payload reference | +| 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 | +| Magic | 8 | ASCII `GTSAMPLE` | +| Trailer checksum | 8 | XXH64 of the preceding 104 trailer bytes, seed zero | + +The trailer is 112 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. + +## Encodings and compression + +Sample column encoding IDs below are distinct from series column encoding IDs, +even where their numeric values coincide. Metadata and compression have separate +namespaces as well. Never reuse an assigned ID with a different meaning. + +| Namespace | ID | Meaning | +| --- | ---: | --- | +| Column encoding | 1 | Timestamp constant step | +| Column encoding | 2 | Timestamp delta-of-delta | +| Column encoding | 3 | Constant sequence | +| Column encoding | 4 | Sequence delta | +| Column encoding | 5 | Nullable Float64 XOR | +| Column encoding | 6 | Constant operation type | +| Column encoding | 7 | Raw operation types | +| Metadata encoding | 1 | Leaf/internal page layout above | +| Compression | 0 | Raw encoded bytes | +| Compression | 1 | Zstd frame | + +Encode first, then compress each chunk/page independently. The default writer +tries Zstd level 3 and chooses it only if strictly smaller than the raw encoded +bytes, following the sample benchmark. Compression 0 requires equal stored and +decoded lengths. Compression level is a writer choice; decoding information is +in the frame. Codec selection adds no public configuration. + +Unsigned varints use base-128/LEB128: seven data bits per byte and a high +continuation bit. Writers emit the shortest representation. Readers reject +truncation, more than ten bytes, a tenth byte greater than 1, and trailing bytes +after the expected values. ZigZag maps signed i64 `x` to +`((x as u64) << 1) ^ ((x >> 63) as u64)` and reverses that mapping on decode. +All wrapping arithmetic below is modulo 2^64, not signed overflowing arithmetic. + +### Timestamp encodings 1 and 2 + +Encoding 1 contains `first: i64` and `step: i64`, exactly 16 bytes. Use it when +every adjacent difference is the same nonnegative, representable i64. A singleton +uses step zero. Reconstruct by checked addition; reject overflow. Zero step is +valid for repeated timestamps. + +Encoding 2 contains the first i64, followed for two or more rows by the first +delta as a ZigZag varint, then ZigZag varints of differences between deltas. +Compute deltas and delta differences with wrapping arithmetic; reconstruct in +the same way. This permits differences across the full signed timestamp range. +After either decoding, require nondecreasing signed timestamps and exact agreement +of first/last values with the descriptor. Consume exactly the recorded row count. +Use encoding 1 when applicable, otherwise encoding 2. + +### Sequence encodings 3 and 4 + +Encoding 3 contains one u64, repeated for all rows. Encoding 4 contains the first +u64 followed by each wrapping u64 difference from its predecessor, reinterpreted +as i64 and written as a ZigZag varint. Decode with wrapping addition. Use 3 if +all sequences are equal, otherwise 4. Preserve all u64 bit patterns. Sequence +need not be monotonic across different timestamps; equal-timestamp runs retain +the input's nonincreasing order. + +### Nullable Float64 encoding 5 + +```text +validity_kind # 0 = all valid, 1 = bitmap +validity # only for kind 1 +first_non_null_bits +xor_deltas +``` + +Bitmap bit `i % 8` in byte `i / 8` is one for a present value; unused high bits of +the final byte are zero. Kind 0 has no bitmap and requires zero null count. Kind +1 requires `0 < null_count < row_count` and an exactly matching count of unset +in-range bits. All-null columns have no descriptor or payload, so every stored +Float64 chunk contains at least one non-null value. + +Process non-null values in row order. Write the first IEEE-754 bit pattern as +u64, then varints of XOR with the previous **non-null** bit pattern. Nulls do not +reset that previous value. Decode exactly `row_count - null_count` values and +scatter according to validity. Preserve all bits, including signed zero, NaN +payloads, infinities, and subnormals. Null and NaN are distinct. The unspecified +value-buffer bits behind nulls are not part of the logical preservation contract. + +### Operation encodings 6 and 7 + +Operation values are u8: 0 delete, 1 put, matching Mito's current `OpType` values. +Unknown operation values fail decoding. Encoding 6 stores one byte repeated for +the whole block; encoding 7 stores exactly one byte per row. Use 6 when all values +are equal, otherwise 7. Neither encoding filters delete rows or discards their +field values. + +### Limits and evolution + +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 +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. + +Unknown file versions, depths, reserved fields, roles, or metadata codecs fail +open/read explicitly. Column role must agree with its encoding: time 1/2, +sequence 3/4, Float64 5, operation 6/7. An unknown column encoding or compression +can be skipped if its chunk is not needed, but attempting to decode it returns +an unsupported-codec error. Known encodings assigned to a different role are +invalid metadata, even for an unprojected column. Unknown fields or malformed +references must not be interpreted as omitted all-null columns. + +New column encodings or compression methods receive new IDs. Structural changes, +changed field meanings, or incompatible limits require a new version. Deeper +metadata trees require an explicit reader-support and rollout decision; an +unchanged layout alone does not make them readable by version-1 implementations. + +## Writer, OpenDAL I/O, and lifecycle + +Implement behind `SamplesFormat::writer` and `FormatWriter`, using the existing +`object_store::ObjectStore` OpenDAL operator. Extend the internal writer factory +with an explicit source SST `FileId` so the manifest does not infer identity by +parsing an opaque prefix. Check that the supplied prefix agrees with metadata, +source ID, and generation. The builder supplies these values from its pinned SST +and captured region version; fake backends follow the same internal interface. + +Validate supported sparse primary-key metadata and extract `(table_id, tsid)` +with the same conventions as the series port. Accept the builder's binary or +dictionary primary-key array representation. Validate column types and mandatory +non-null values before writing; preserve ordering state across `write` calls. +Retain only the active block and bounded per-part descriptors, not all SST rows. +No external sort, benchmark manifest/schema file, or filesystem scratch spool is +introduced. + +Write chunks as each block completes. Finish a part at the next block boundary +once its stored position reaches 512 MiB or its accumulated serialized block +descriptors reach 8 MiB. Both are soft packing limits; one bounded block may +overshoot. Root capacity is a separate hard limit checked before emitting that +block. Write leaves, the internal root if needed, and the trailer, then close the +part before opening its successor. Bound the manifest and lifecycle object-list +growth as parts complete; fail rather than constructing a manifest over 64 MiB. + +Use recorded sizes for reads; do not issue object stats just to locate trailers. +Fetch metadata and selected column ranges only. Reads are sequential in this +initial port, without speculative adjacent-field fetches or prefetch. Keep +fetch/validate/decode entrypoints for metadata and columns separate so a later +cache integration can use immutable generation, part, byte range, checksum, +codec IDs, and representation without changing the format. + +Writers own all partial outputs until `finish` succeeds. Write/close errors, +explicit abort, and cancellation clean up every created object, including an +unfinished manifest. Use the existing cleanup/runtime pattern from #9841; +cleanup errors must not hide the original error. A process crash can leave orphan +objects, but never publishes an unfinished generation as valid. Reconciliation +and existing orphan handling remain responsible for recovery. + +After finish, ownership transfers to the build/catalog lifecycle. Failed catalog +publication retires unpublished outputs and preserves the prior catalog. Each +reader and returned stream holds a `SamplesFileHandle`; superseded objects are +purged only after those handles are released. Sample sidecars remain rebuildable, +and authoritative data stays in Parquet SSTs. + +## Reader API and execution + +Add a sample-specific reader factory to `SamplesFormat`, accepting the store and +a pinned `SamplesFileHandle`. Provide access to the handle's region/source identity +for manifest validation. Return a `SamplesFormatReader` exposing `schema()` and +`read(SamplesReadRequest)` as an independent stream of projected Arrow batches. +Do not reuse the series-only read request: it has no time semantics. + +The internal request has these fields: + +| Field | Meaning | +| --- | --- | +| `projection: Vec` | Ordered names from the stored logical schema | +| `table_id: Option` | Optional table restriction | +| `ids: Option>` | Optional qualified identity set | +| `start: Option` | Inclusive timestamp in the stored unit | +| `end: Option` | Exclusive timestamp in the stored unit | + +`ids = None` means unrestricted; an explicitly empty set means zero rows. +Intersect `ids` with `table_id` when both are supplied. Validate unknown or +duplicate projection names before part I/O; preserve requested output order. +An empty projection preserves selected row counts. Equal time endpoints produce +an empty stream, and a reversed interval is an error. Unbounded endpoints avoid +sentinel arithmetic and permit selecting `i64::MAX` with no end bound. Callers +perform any unit conversion explicitly; the format API does not guess units. + +1. The factory opens the bounded manifest, validates handle identity and recorded + sizes, and exposes the stored schema. Each `read` validates its request even + if the artifact or selection is empty. An empty selection needs no part reads. +2. Prune parts using inclusive ID ranges and time envelopes. Read and validate + selected headers/trailers, then traverse only candidate pages in physical + order. At depth 2, retain the root references and one active leaf. +3. Select blocks by exact identity and overlapping time bounds. Decode timestamp + chunks for a constrained time request, validate their bounds, and select rows. + With no time restriction, an unprojected timestamp chunk need not be fetched. +4. Fetch the remaining projected chunks only for blocks with surviving rows, + reusing a timestamp chunk already decoded for selection. Decode each chunk + once per block, synthesize identity constants and omitted null fields, and + emit selected rows in input order. One output batch per selected block is the + initial packing choice, with the additional identity-column accounting above. + +Sequence and operation chunks are mandatory on disk but are read only if +projected. The format reader does not silently apply visibility or delete rules. +A future query adapter must request them when required for merging. There are +no tag/value predicates in this initial request, no selector-result cache, and +no cross-read page/data cache. All reads have independent traversal state and +can outlive the factory. Validation is incremental over visited metadata and +decoded chunks, not an eager whole-artifact integrity scan; unselected payload +corruption is not guaranteed to be discovered by a selective query. +Check block ordinals and row ordering across adjacent blocks when both are read; +do not require the first selected block to have ordinal zero after pruning. +A complete traversal also verifies the manifest's row and block totals. These +totals are diagnostic summaries, not allocation sizes or a replacement for block +row counts when producing output. + +## Compatibility, validation, and deferred work + +The backend is opt-in through internal format selection, like #9841. Existing catalogs, Parquet +artifacts, and series artifacts retain their meaning. An older binary is not +promised to understand the new catalog identifier or sample bytes; do not claim +downgrade compatibility merely because the samples are rebuildable. The new +backend must not become necessary for reading authoritative SSTs. Query-consumer +integration and a public enablement/rollout mechanism are separate work. + +Implementation acceptance includes: + +- Codec tests and round trips for every supported timestamp unit, repeated and + negative timestamps, signed extremes, full-range sequences, both operations, + null/NaN distinction, Float64 bit patterns, and nulls between XOR values. +- Empty input, singleton blocks, all-null/missing fields, mixed tables and sparse + schemas, identical TSIDs in different tables, cross-batch continuation, + duplicate/conflicting rows, and decreasing-input rejection. +- Row and decoded-byte caps with narrow/wide schemas; smaller row limits; blocks + spanning hours; forced payload/metadata rollover; depths 1/2 and root-capacity + rollover without creating depth 3. Reject depths 0/3 and internal children. +- Projection order and empty projection, all/empty/absent ID sets, table/ID + intersection, exact time boundaries, unbounded endpoints, and equal timestamps + spanning multiple leaves and parts. Verify selected physical ranges: no unused + value chunks, duplicate page reads, or hidden stat requests. +- Truncation, corrupt checksums, invalid area bounds, u64 offset overflow, + references beyond 4 GiB, wrong lengths/counts/roles/ordinals, inconsistent + time bounds, malformed validity/varints, unsupported codecs, and bounded + decompression. Test malformed metadata without allocating from untrusted counts. +- Manifest/schema/handle identity mismatches, wrong or duplicate object paths, + size accounting, reopen and reconciliation, OpenDAL write/close failures, + abort/cancellation, failed catalog publication, and retirement with pinned + streams on memory and filesystem stores. +- A persisted-format compatibility case under `tests/compatibility/`, following + its [README](../../tests/compatibility/README.md) and + [agent guide](../../tests/compatibility/AGENTS.md), covering the approved version-1 + bytes and legacy catalog behavior. Do not substitute a SQL query test that + never enables this internal backend or claim cross-binary coverage from a + same-build round trip. +- `cargo nextest run -p mito2`, crate Clippy, formatting, and license checks for + new sources. Any broader workspace validation and its limitations must be + reported when the implementation PR is prepared. + +The benchmark comparison changes materially: nullable fields and operation types +add bytes; native timestamp units and preserved duplicate rows change codec +inputs; 64-bit checksums, Mito manifests/references, and overlapping page bounds +change framing; bounded two-level trees replace disk-spooled arbitrary-depth +indexes; row/byte packing replaces the benchmark's hourly default. OpenDAL and +the absence of caches also change performance. Measure the port independently; +no performance equivalence is asserted. + +The benchmark follow-ups from [PR #9887](https://github.com/GreptimeTeam/greptimedb/pull/9887) +include: + +- High-churn series with only a few rows per SST. Report stored bytes per row + and separate metadata from column payloads across different rows-per-series + distributions. Each identity change starts a new block, whose descriptor is + 220 bytes before metadata compression with one present field, in addition to + the mandatory timestamp, sequence, and operation chunks and the value chunk. +- Integer-valued Float64 counters. Compare the current XOR encoding with an + exact `first`/`step` progression encoding like the prototype in + [PR #9556](https://github.com/GreptimeTeam/greptimedb/pull/9556), measuring encoded + and compressed sizes as well as encode/decode throughput. Include steady + counters, resets, irregular increments, and nulls; any candidate must preserve + Float64 bits exactly and fall back for values it cannot represent losslessly. + +These measurements are deferred and do not change the version-1 wire contract. + +Additional field types, value statistics/predicates, exact primary keys, +dedicated job/instance identifier components, label reconstruction, cache +integration, deeper trees, and query integration are deferred. The byte assignments +above define the version-1 contract. The +[frozen sample fixture](../../tests/compatibility/wire/mito_columnar_samples/README.md) +checks those bytes independently of writer/reader round trips. + +[bench-spec]: https://github.com/GreptimeTeam/series-format-bench/blob/8650c3003a861c8cf2c9f5037bfc81fde6a3bdc4/docs/formats/columnar-samples-v2.md +[bench-code]: https://github.com/GreptimeTeam/series-format-bench/blob/8650c3003a861c8cf2c9f5037bfc81fde6a3bdc4/src/sample_columnar.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index be100512f10..421b0df6723 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -69,6 +69,11 @@ filtered `RecordBatch` stream. - **Series format PoC**: samples are rebuildable local SST sidecars. Optional writer interfaces (`format.rs`) integrate with reconciliation, catalog publication, and deferred purge. Defaults retain the original Parquet writer and disable samples. The opt-in columnar backend (`series_index/columnar.rs`) implements streaming OpenDAL writes and pinned format reads; its wire contract is in `docs/formats/mito-series.md`. Artifact objects record paths and sizes, including the manifest, so reads need no stat. Codecs currently run inline pending scheduling benchmarks. Prepared sources and query-consumer integration are deferred; ordinary scans only use Parquet index artifacts. - **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. +- **Columnar samples**: `series_index/samples.rs` implements opt-in per-SST sidecars, + including nullable Float64 fields and preserved sequence/delete rows. Its wire + contract is `docs/formats/mito-samples.md`; time/ID bounds may overlap across + leaves and parts. Both columnar backends share cancellation-safe output ownership + in `series_index/artifact_writer.rs`. Sample reads do not apply query visibility. - **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`. - **WAL/memtable encoding** (`wal/`, `memtable/`): breaks replay if changed incompatibly. diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 57b3611f999..dc8b461afd6 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -17,6 +17,7 @@ //! Under development. Index files are currently stored on the local filesystem. pub(crate) mod artifact; +mod artifact_writer; mod bucket; mod builder; mod catalog; @@ -30,6 +31,11 @@ pub(crate) mod format; mod format_tests; mod maintenance; mod purger; +#[allow( + dead_code, + reason = "Samples are an opt-in backend; query integration is deferred." +)] +pub(crate) mod samples; mod searcher; mod task; #[cfg(test)] diff --git a/src/mito2/src/series_index/artifact.rs b/src/mito2/src/series_index/artifact.rs index fdce49b31fd..b5db30b3a58 100644 --- a/src/mito2/src/series_index/artifact.rs +++ b/src/mito2/src/series_index/artifact.rs @@ -30,6 +30,8 @@ pub(crate) enum ArtifactFormat { Parquet, #[serde(rename = "columnar-series")] ColumnarSeries, + #[serde(rename = "columnar-samples")] + ColumnarSamples, #[cfg(test)] #[serde(rename = "fake-samples")] FakeSamples, @@ -83,6 +85,10 @@ impl SamplesFileHandle { })) } + pub(crate) fn file_id(&self) -> RegionFileId { + self.0.file_id + } + pub(crate) fn entry(&self) -> &SamplesEntry { &self.0.entry } diff --git a/src/mito2/src/series_index/artifact_writer.rs b/src/mito2/src/series_index/artifact_writer.rs new file mode 100644 index 00000000000..977e74efc15 --- /dev/null +++ b/src/mito2/src/series_index/artifact_writer.rs @@ -0,0 +1,144 @@ +// 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. + +//! Owned outputs shared by immutable columnar artifacts. + +use bytes::Bytes; +use object_store::{ObjectStore, Writer}; +use snafu::{IntoError, ResultExt}; + +use crate::error::{JoinSnafu, OpenDalSnafu, Result}; +use crate::series_index::columnar::{invalid, require}; +use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; + +/// Drop owns the live OpenDAL writer, including its private atomic-write temporary file. +pub(super) struct LiveFile { + writer: Option, + store: ObjectStore, + path: String, +} +impl Drop for LiveFile { + fn drop(&mut self) { + if let Some(mut writer) = self.writer.take() { + let store = self.store.clone(); + let path = self.path.clone(); + common_runtime::spawn_compact(async move { + if let Err(error) = writer.abort().await { + common_telemetry::warn!(error; "Aborting artifact output {path}"); + } + if let Err(error) = store.delete(&path).await { + common_telemetry::warn!(error; "Deleting artifact output {path}"); + } + }); + } + } +} + +pub(super) struct Outputs { + pub(super) store: ObjectStore, + pub(super) active: Option, + pub(super) paths: Vec, + pub(super) committed: bool, +} +impl Outputs { + pub(super) async fn open(&mut self, path: String) -> Result<()> { + require(self.active.is_none(), "artifact output already open")?; + self.paths.push(path.clone()); + let store = self.store.clone(); + // The task owns opening: cancellation drops the returned LiveFile and aborts it. + self.active = Some( + common_runtime::spawn_compact(async move { + let writer = store + .writer_with(&path) + .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .concurrent(DEFAULT_WRITE_CONCURRENCY) + .await + .context(OpenDalSnafu)?; + Ok::<_, crate::error::Error>(LiveFile { + writer: Some(writer), + store, + path, + }) + }) + .await + .context(JoinSnafu)??, + ); + Ok(()) + } + pub(super) async fn write(&mut self, bytes: Vec) -> Result<()> { + self.active + .as_mut() + .and_then(|f| f.writer.as_mut()) + .ok_or_else(|| invalid("artifact output is not open"))? + .write(Bytes::from(bytes)) + .await + .context(OpenDalSnafu)?; + Ok(()) + } + pub(super) async fn close(&mut self) -> Result<()> { + let file = self + .active + .as_mut() + .ok_or_else(|| invalid("artifact output is not open"))?; + file.writer + .as_mut() + .ok_or_else(|| invalid("artifact output is closed"))? + .close() + .await + .context(OpenDalSnafu)?; + file.writer = None; + self.active = None; + Ok(()) + } + pub(super) async fn cleanup(&mut self) -> Result<()> { + let mut first_error = None; + if let Some(file) = self.active.as_mut() + && let Some(writer) = file.writer.as_mut() + { + match writer.abort().await { + Ok(()) => file.writer = None, + Err(error) => first_error = Some(OpenDalSnafu.into_error(error)), + } + } + self.active = None; + for path in &self.paths { + if let Err(error) = self.store.delete(path).await { + first_error.get_or_insert_with(|| OpenDalSnafu.into_error(error)); + } + } + if let Some(error) = first_error { + Err(error) + } else { + self.committed = true; + Ok(()) + } + } +} +impl Drop for Outputs { + fn drop(&mut self) { + if !self.committed { + let mut owned = Self { + store: self.store.clone(), + active: self.active.take(), + paths: std::mem::take(&mut self.paths), + committed: true, + }; + common_runtime::spawn_compact(async move { + if let Err(error) = owned.cleanup().await { + common_telemetry::warn!(error; "Cleaning unpublished format artifact"); + } + }); + } + } +} diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index 2a9db9503e3..c875456a819 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -277,7 +277,13 @@ pub(crate) async fn build_samples( file_id.file_id() ); let mut writer = backend - .writer(version.metadata.clone(), store.clone(), &prefix, generation) + .writer( + version.metadata.clone(), + store.clone(), + &prefix, + file_id.file_id(), + generation, + ) .await?; let result: Result<()> = async { while let Some((row_group, selection)) = selection.pop_first() { diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index 24c8fb27086..6cb7b33912b 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -314,6 +314,7 @@ mod tests { for (format, name) in [ (ArtifactFormat::Parquet, "parquet"), (ArtifactFormat::ColumnarSeries, "columnar-series"), + (ArtifactFormat::ColumnarSamples, "columnar-samples"), (ArtifactFormat::FakeSamples, "fake-samples"), (ArtifactFormat::FakeIndex, "fake-index"), ] { diff --git a/src/mito2/src/series_index/columnar.rs b/src/mito2/src/series_index/columnar.rs index d010e7356fd..5ebe8d86fcf 100644 --- a/src/mito2/src/series_index/columnar.rs +++ b/src/mito2/src/series_index/columnar.rs @@ -14,7 +14,7 @@ //! Columnar series format. See `docs/formats/mito-series.md` for the wire contract. -mod codec; +pub(super) mod codec; mod metadata; mod reader; #[cfg(test)] diff --git a/src/mito2/src/series_index/columnar/writer.rs b/src/mito2/src/series_index/columnar/writer.rs index 72156ff139c..05fe94ee81b 100644 --- a/src/mito2/src/series_index/columnar/writer.rs +++ b/src/mito2/src/series_index/columnar/writer.rs @@ -15,17 +15,17 @@ //! Streaming aggregation and owned OpenDAL outputs. use async_trait::async_trait; -use bytes::Bytes; use datatypes::arrow::array::{Array, BinaryArray, DictionaryArray}; use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; use mito_codec::row_converter::SparseOffsetsCache; -use object_store::{ObjectStore, Writer}; -use snafu::{IntoError, ResultExt}; +use object_store::ObjectStore; +use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; -use crate::error::{JoinSnafu, OpenDalSnafu, Result, SerdeJsonSnafu}; +use crate::error::{Result, SerdeJsonSnafu}; use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata, ArtifactObject}; +use crate::series_index::artifact_writer::Outputs; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::columnar::codec::{self, encode_dictionary, encode_integers}; use crate::series_index::columnar::metadata::{ @@ -37,132 +37,10 @@ use crate::series_index::columnar::{ use crate::series_index::format::FormatWriter; use crate::series_index::writer::{decode_series_identity, tag_columns, validate_metadata}; use crate::series_index::{MetricSeriesId, ROW_COUNT_COLUMN, TSID_COLUMN}; -use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; /// Maximum owned row bytes per block; a single larger row is flushed on its own. const BLOCK_BYTE_BUDGET: usize = 8 * 1024 * 1024; -/// Drop owns the live OpenDAL writer, including its private atomic-write temporary file. -struct LiveFile { - writer: Option, - store: ObjectStore, - path: String, -} -impl Drop for LiveFile { - fn drop(&mut self) { - if let Some(mut writer) = self.writer.take() { - let store = self.store.clone(); - let path = self.path.clone(); - common_runtime::spawn_compact(async move { - if let Err(error) = writer.abort().await { - common_telemetry::warn!(error; "Aborting series output {path}"); - } - if let Err(error) = store.delete(&path).await { - common_telemetry::warn!(error; "Deleting series output {path}"); - } - }); - } - } -} - -struct Outputs { - store: ObjectStore, - active: Option, - paths: Vec, - committed: bool, -} -impl Outputs { - async fn open(&mut self, path: String) -> Result<()> { - require(self.active.is_none(), "series output already open")?; - self.paths.push(path.clone()); - let store = self.store.clone(); - // The task owns opening: cancellation drops the returned LiveFile and aborts it. - self.active = Some( - common_runtime::spawn_compact(async move { - let writer = store - .writer_with(&path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) - .concurrent(DEFAULT_WRITE_CONCURRENCY) - .await - .context(OpenDalSnafu)?; - Ok::<_, crate::error::Error>(LiveFile { - writer: Some(writer), - store, - path, - }) - }) - .await - .context(JoinSnafu)??, - ); - Ok(()) - } - async fn write(&mut self, bytes: Vec) -> Result<()> { - self.active - .as_mut() - .and_then(|f| f.writer.as_mut()) - .ok_or_else(|| invalid("series output is not open"))? - .write(Bytes::from(bytes)) - .await - .context(OpenDalSnafu)?; - Ok(()) - } - async fn close(&mut self) -> Result<()> { - let file = self - .active - .as_mut() - .ok_or_else(|| invalid("series output is not open"))?; - file.writer - .as_mut() - .ok_or_else(|| invalid("series output is closed"))? - .close() - .await - .context(OpenDalSnafu)?; - file.writer = None; - self.active = None; - Ok(()) - } - async fn cleanup(&mut self) -> Result<()> { - let mut first_error = None; - if let Some(file) = self.active.as_mut() - && let Some(writer) = file.writer.as_mut() - { - match writer.abort().await { - Ok(()) => file.writer = None, - Err(error) => first_error = Some(OpenDalSnafu.into_error(error)), - } - } - self.active = None; - for path in &self.paths { - if let Err(error) = self.store.delete(path).await { - first_error.get_or_insert_with(|| OpenDalSnafu.into_error(error)); - } - } - if let Some(error) = first_error { - Err(error) - } else { - self.committed = true; - Ok(()) - } - } -} -impl Drop for Outputs { - fn drop(&mut self) { - if !self.committed { - let mut owned = Self { - store: self.store.clone(), - active: self.active.take(), - paths: std::mem::take(&mut self.paths), - committed: true, - }; - common_runtime::spawn_compact(async move { - if let Err(error) = owned.cleanup().await { - common_telemetry::warn!(error; "Cleaning unpublished series artifact"); - } - }); - } - } -} - #[derive(Debug)] struct Row { id: MetricSeriesId, diff --git a/src/mito2/src/series_index/format.rs b/src/mito2/src/series_index/format.rs index 43a8ff39d26..eb552a8eabe 100644 --- a/src/mito2/src/series_index/format.rs +++ b/src/mito2/src/series_index/format.rs @@ -28,7 +28,7 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::FileId; use crate::error::Result; -use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata}; +use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata, SamplesFileHandle}; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::{MetricSeriesId, SeriesIndexFileHandle}; @@ -87,16 +87,45 @@ pub(crate) trait SeriesIndexFormat: Debug + Send + Sync { ) -> Result>; } -/// Builds samples artifacts; read interfaces are deferred. +/// Selection in the stored timestamp unit. No visibility filtering is implicit. +#[allow(dead_code, reason = "Query adapter is deferred.")] +#[derive(Debug, Clone, Default)] +pub(crate) struct SamplesReadRequest { + pub(crate) projection: Vec, + pub(crate) table_id: Option, + pub(crate) ids: Option>, + pub(crate) start: Option, + pub(crate) end: Option, +} + +#[allow(dead_code, reason = "Query adapter is deferred.")] +pub(crate) trait SamplesFormatReader: Send + Sync { + fn schema(&self) -> SchemaRef; + fn read(&self, request: SamplesReadRequest) -> Result>>; +} + +/// Builds and opens immutable per-SST samples artifacts. #[async_trait] pub(crate) trait SamplesFormat: Debug + Send + Sync { fn format(&self) -> ArtifactFormat; + #[allow(dead_code, reason = "Query adapter is deferred.")] + async fn reader( + &self, + _store: ObjectStore, + _handle: SamplesFileHandle, + ) -> Result> { + crate::error::UnsupportedOperationSnafu { + err_msg: "This samples format has no reader", + } + .fail() + } /// The caller supplies a unique prefix. Finish reports all objects owned below it. async fn writer( &self, metadata: RegionMetadataRef, store: ObjectStore, prefix: &str, + source_file_id: FileId, generation: FileId, ) -> Result>; } diff --git a/src/mito2/src/series_index/format_tests.rs b/src/mito2/src/series_index/format_tests.rs index 0fa2661c93e..e8fba18de92 100644 --- a/src/mito2/src/series_index/format_tests.rs +++ b/src/mito2/src/series_index/format_tests.rs @@ -84,6 +84,7 @@ impl SamplesFormat for FakeSamples { _metadata: RegionMetadataRef, store: ObjectStore, prefix: &str, + _source_file_id: FileId, generation: FileId, ) -> Result> { let objects = vec![format!("{prefix}data"), format!("{prefix}directory")]; diff --git a/src/mito2/src/series_index/samples.rs b/src/mito2/src/series_index/samples.rs new file mode 100644 index 00000000000..23d34f02900 --- /dev/null +++ b/src/mito2/src/series_index/samples.rs @@ -0,0 +1,222 @@ +// 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. + +//! Columnar samples sidecars. Wire contract: `docs/formats/mito-samples.md`. + +mod codec; +mod metadata; +mod reader; +#[cfg(test)] +mod tests; +mod writer; + +use std::collections::HashSet; +use std::sync::Arc; + +use async_trait::async_trait; +use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use object_store::ObjectStore; +use serde::{Deserialize, Serialize}; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::FileId; + +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::{MetricSeriesId, TABLE_ID_COLUMN, TSID_COLUMN}; + +const VERSION: u32 = 1; +const MANIFEST_LIMIT: usize = 64 * 1024 * 1024; +const MAX_ROWS: usize = 8192; +const MAX_FIELDS: usize = 1024; +const BLOCK_BYTES: usize = 256 * 1024; +const SEQUENCE: &str = "__sequence"; +const OPERATION: &str = "__op_type"; + +/// Internal packing knobs; the backend is disabled unless selected explicitly. +#[derive(Debug, Clone)] +pub(crate) struct ColumnarSamplesFormat { + pub(crate) block_rows: usize, + pub(crate) part_bytes: u64, + pub(crate) metadata_bytes: usize, + pub(crate) page_bytes: usize, + pub(crate) root_capacity: usize, +} +impl Default for ColumnarSamplesFormat { + fn default() -> Self { + Self { + block_rows: MAX_ROWS, + part_bytes: 512 * 1024 * 1024, + metadata_bytes: 8 * 1024 * 1024, + page_bytes: metadata::PAGE_LIMIT, + root_capacity: metadata::ROOT_CAPACITY, + } + } +} +#[async_trait] +impl SamplesFormat for ColumnarSamplesFormat { + fn format(&self) -> ArtifactFormat { + ArtifactFormat::ColumnarSamples + } + async fn writer( + &self, + metadata: RegionMetadataRef, + store: ObjectStore, + prefix: &str, + source_file_id: FileId, + generation: FileId, + ) -> Result> { + Ok(Box::new(writer::SamplesWriter::new( + self.clone(), + metadata, + store, + prefix, + source_file_id, + generation, + )?)) + } + async fn reader( + &self, + store: ObjectStore, + handle: SamplesFileHandle, + ) -> Result> { + Ok(Box::new(reader::SamplesReader::open(store, handle).await?)) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct TimestampColumn { + column_id: u32, + name: String, + unit: TimestampUnit, +} +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +enum TimestampUnit { + Second, + Millisecond, + Microsecond, + Nanosecond, +} +impl TimestampUnit { + fn arrow(self) -> TimeUnit { + match self { + Self::Second => TimeUnit::Second, + Self::Millisecond => TimeUnit::Millisecond, + Self::Microsecond => TimeUnit::Microsecond, + Self::Nanosecond => TimeUnit::Nanosecond, + } + } + fn from_arrow(unit: &TimeUnit) -> Self { + match unit { + TimeUnit::Second => Self::Second, + TimeUnit::Millisecond => Self::Millisecond, + TimeUnit::Microsecond => Self::Microsecond, + TimeUnit::Nanosecond => Self::Nanosecond, + } + } +} +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ValueColumn { + column_id: u32, + name: String, + #[serde(rename = "type")] + data_type: String, + nullable: bool, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Part { + part_id: u32, + file_size: u64, + num_rows: u64, + num_blocks: u64, + first: MetricSeriesId, + last: MetricSeriesId, + min_timestamp: i64, + max_timestamp: i64, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Manifest { + format: ArtifactFormat, + version: u32, + region_id: u64, + source_file_id: FileId, + generation: FileId, + timestamp: TimestampColumn, + fields: Vec, + files: Vec, +} +impl Manifest { + fn prefix(&self) -> String { + format!( + "{}/samples/{}/{}/", + self.region_id, self.source_file_id, self.generation + ) + } + fn path(&self) -> String { + format!("{}manifest.json", self.prefix()) + } + fn part_path(&self, id: u32) -> String { + format!("{}part-{id}.samples", self.prefix()) + } + fn validate_schema(&self) -> Result<()> { + require( + (1..=MAX_FIELDS).contains(&self.fields.len()), + "invalid samples field count", + )?; + let mut names = HashSet::from([TABLE_ID_COLUMN, TSID_COLUMN, SEQUENCE, OPERATION]); + require( + names.insert(&self.timestamp.name), + "duplicate samples timestamp name", + )?; + let mut ids = HashSet::from([self.timestamp.column_id]); + for field in &self.fields { + require( + field.data_type == "float64" && field.nullable, + "unsupported samples field type/nullability", + )?; + require( + names.insert(&field.name) && ids.insert(field.column_id), + "duplicate samples field name/id", + )?; + } + Ok(()) + } + fn schema(&self) -> SchemaRef { + let mut fields = vec![ + Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), + Field::new(TSID_COLUMN, DataType::UInt64, false), + Field::new( + &self.timestamp.name, + DataType::Timestamp(self.timestamp.unit.arrow(), None), + false, + ), + Field::new(SEQUENCE, DataType::UInt64, false), + Field::new(OPERATION, DataType::UInt8, false), + ]; + fields.extend( + self.fields + .iter() + .map(|f| Field::new(&f.name, DataType::Float64, true)), + ); + Arc::new(Schema::new(fields)) + } +} + +fn valid_rows(rows: usize, fields: usize) -> bool { + (1..=MAX_ROWS).contains(&rows) + && (1..=MAX_FIELDS).contains(&fields) + && rows * (17 + 8 * fields) + fields * rows.div_ceil(8) <= BLOCK_BYTES +} diff --git a/src/mito2/src/series_index/samples/codec.rs b/src/mito2/src/series_index/samples/codec.rs new file mode 100644 index 00000000000..145c545c647 --- /dev/null +++ b/src/mito2/src/series_index/samples/codec.rs @@ -0,0 +1,225 @@ +// 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. + +//! Lossless, bounded sample column encodings. Nulls never become NaNs. + +use crate::error::Result; +pub(super) use crate::series_index::columnar::codec::{ + Cursor, checksum, decompress, put_u32, put_u64, put_varint, +}; +use crate::series_index::samples::{MAX_ROWS, require, unsupported}; + +fn zig(x: i64) -> u64 { + ((x as u64) << 1) ^ ((x >> 63) as u64) +} +fn unzig(x: u64) -> u64 { + (x >> 1) ^ (0u64.wrapping_sub(x & 1)) +} + +pub(super) fn encode_time(values: &[i64]) -> (u16, Vec) { + let mut out = values[0].to_le_bytes().to_vec(); + let step = if values.len() == 1 { + Some(0) + } else { + values[1].checked_sub(values[0]) + }; + if let Some(step) = step.filter(|&step| { + step >= 0 + && values + .windows(2) + .all(|v| v[1].checked_sub(v[0]) == Some(step)) + }) { + put_u64(&mut out, step as u64); + (1, out) + } else { + let mut prior = 0u64; + for (i, pair) in values.windows(2).enumerate() { + let delta = (pair[1] as u64).wrapping_sub(pair[0] as u64); + put_varint( + &mut out, + zig(if i == 0 { + delta + } else { + delta.wrapping_sub(prior) + } as i64), + ); + prior = delta; + } + (2, out) + } +} +pub(super) fn encode_sequence(values: &[u64]) -> (u16, Vec) { + let mut out = values[0].to_le_bytes().to_vec(); + if values.iter().all(|v| *v == values[0]) { + return (3, out); + } + for pair in values.windows(2) { + put_varint(&mut out, zig(pair[1].wrapping_sub(pair[0]) as i64)); + } + (4, out) +} +pub(super) fn decode_numbers(bytes: &[u8], encoding: u16, rows: usize) -> Result> { + require( + (1..=MAX_ROWS).contains(&rows), + "invalid samples number count", + )?; + if !(1..=4).contains(&encoding) { + return Err(unsupported("samples number encoding")); + } + let mut c = Cursor::new(bytes); + let first = c.u64()?; + let mut values = Vec::with_capacity(rows); + values.push(first); + if encoding == 1 { + let step = c.u64()? as i64; + require( + step >= 0 && (rows != 1 || step == 0), + "invalid sample timestamp step", + )?; + let mut last = first as i64; + for _ in 1..rows { + last = last + .checked_add(step) + .ok_or_else(|| super::invalid("sample timestamp overflow"))?; + values.push(last as u64); + } + } else if encoding == 3 { + values.resize(rows, first); + } else { + require( + rows - 1 <= c.remaining(), + "sample number count exceeds bytes", + )?; + let mut previous = first; + let mut delta = 0u64; + for i in 1..rows { + let next = unzig(c.varint()?); + delta = if encoding == 4 || i == 1 { + next + } else { + delta.wrapping_add(next) + }; + previous = previous.wrapping_add(delta); + values.push(previous); + } + } + c.end()?; + if encoding <= 2 { + require( + values.windows(2).all(|v| (v[0] as i64) <= v[1] as i64), + "unordered sample timestamps", + )?; + } + Ok(values) +} + +pub(super) fn encode_float(values: &[Option]) -> Vec { + let nullable = values.iter().any(Option::is_none); + let mut out = vec![u8::from(nullable)]; + if nullable { + out.resize(1 + values.len().div_ceil(8), 0); + for (i, value) in values.iter().enumerate() { + if value.is_some() { + out[1 + i / 8] |= 1 << (i % 8); + } + } + } + let mut previous = None; + for &value in values.iter().flatten() { + if let Some(last) = previous { + put_varint(&mut out, value ^ last); + } else { + put_u64(&mut out, value); + } + previous = Some(value); + } + out +} +pub(super) fn decode_float(bytes: &[u8], rows: usize, nulls: u32) -> Result>> { + require( + (1..=MAX_ROWS).contains(&rows) && (nulls as usize) < rows, + "invalid samples float count", + )?; + let mut c = Cursor::new(bytes); + let bitmap = match c.u8()? { + 0 => { + require(nulls == 0, "samples float null count mismatch")?; + None + } + 1 => { + require(nulls > 0, "redundant samples validity bitmap")?; + let bitmap = c.take(rows.div_ceil(8))?; + if !rows.is_multiple_of(8) { + require( + bitmap[bitmap.len() - 1] >> (rows % 8) == 0, + "nonzero sample validity padding", + )?; + } + let present: u32 = bitmap.iter().map(|b| b.count_ones()).sum(); + require( + present as usize + nulls as usize == rows, + "samples validity count mismatch", + )?; + Some(bitmap) + } + _ => return Err(unsupported("samples validity kind")), + }; + let first = c.u64()?; + require( + rows - nulls as usize - 1 <= c.remaining(), + "sample float count exceeds bytes", + )?; + let mut previous = None; + let mut out = Vec::with_capacity(rows); + for i in 0..rows { + if bitmap.is_none_or(|b| b[i / 8] & (1 << (i % 8)) != 0) { + let value = match previous { + None => first, + Some(last) => last ^ c.varint()?, + }; + out.push(Some(value)); + previous = Some(value); + } else { + out.push(None); + } + } + c.end()?; + Ok(out) +} +pub(super) fn encode_operations(values: &[u8]) -> (u16, Vec) { + if values.iter().all(|v| *v == values[0]) { + (6, vec![values[0]]) + } else { + (7, values.to_vec()) + } +} +pub(super) fn decode_operations(bytes: &[u8], encoding: u16, rows: usize) -> Result> { + require( + (1..=MAX_ROWS).contains(&rows), + "invalid samples operation count", + )?; + let values = match encoding { + 6 => { + require(bytes.len() == 1, "invalid constant sample operation")?; + vec![bytes[0]; rows] + } + 7 => { + require(bytes.len() == rows, "invalid sample operation length")?; + bytes.to_vec() + } + _ => return Err(unsupported("samples operation encoding")), + }; + require(values.iter().all(|v| *v <= 1), "unknown sample operation")?; + Ok(values) +} diff --git a/src/mito2/src/series_index/samples/metadata.rs b/src/mito2/src/series_index/samples/metadata.rs new file mode 100644 index 00000000000..a3d1c910106 --- /dev/null +++ b/src/mito2/src/series_index/samples/metadata.rs @@ -0,0 +1,500 @@ +// 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. + +//! Explicit sample framing and bounded two-level metadata pages. + +use crate::error::Result; +use crate::series_index::samples::codec::{self, Cursor, put_u32, put_u64}; +use crate::series_index::samples::{ + Manifest, MetricSeriesId, VERSION, invalid, require, unsupported, valid_rows, +}; + +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 PAGE_LIMIT: usize = 65536; +pub(super) const ROOT_CAPACITY: usize = (PAGE_LIMIT - 12) / 72; +const CHUNK_LIMIT: usize = 1024 * 1024; + +#[derive(Debug, Clone, Copy)] +pub(super) struct PayloadReference { + pub offset: u64, + pub stored: u32, + pub decoded: u32, + pub checksum: u64, + pub encoding: u16, + pub compression: u16, +} +impl PayloadReference { + pub fn end(&self) -> Result { + self.offset + .checked_add(u64::from(self.stored)) + .ok_or_else(|| invalid("sample payload offset overflow")) + } + pub fn validate(&self, start: u64, end: u64, page: bool) -> Result<()> { + let limit = if page { PAGE_LIMIT } else { CHUNK_LIMIT }; + require( + self.stored > 0 + && self.decoded > 0 + && self.decoded as usize <= limit + && self.stored as usize <= limit + 65536, + "sample payload exceeds limits", + )?; + require( + self.offset >= start && self.end()? <= end, + "sample payload outside area", + )?; + if self.compression == 0 { + require(self.stored == self.decoded, "raw samples length mismatch")?; + } + Ok(()) + } + fn encode(&self, out: &mut Vec) { + put_u64(out, self.offset); + put_u32(out, self.stored); + put_u32(out, self.decoded); + put_u64(out, self.checksum); + out.extend(self.encoding.to_le_bytes()); + out.extend(self.compression.to_le_bytes()); + put_u32(out, 0); + } + fn decode(c: &mut Cursor<'_>) -> Result { + let result = Self { + offset: c.u64()?, + stored: c.u32()?, + decoded: c.u32()?, + checksum: c.u64()?, + encoding: c.u16()?, + compression: c.u16()?, + }; + require(c.u32()? == 0, "nonzero sample payload reserved field")?; + Ok(result) + } + pub fn unpack(&self, bytes: &[u8]) -> Result> { + require( + bytes.len() == self.stored as usize && codec::checksum(bytes) == self.checksum, + "sample checksum mismatch", + )?; + codec::decompress(bytes, self.compression, self.decoded as usize) + } +} +pub(super) fn pack( + bytes: &[u8], + offset: u64, + encoding: u16, + page: bool, +) -> Result<(PayloadReference, Vec)> { + require( + !bytes.is_empty() && bytes.len() <= if page { PAGE_LIMIT } else { CHUNK_LIMIT }, + "sample encoded payload exceeds limit", + )?; + let compressed = + zstd::bulk::compress(bytes, 3).map_err(|e| invalid(format!("sample compression: {e}")))?; + let (compression, stored) = if compressed.len() < bytes.len() { + (1, compressed) + } else { + (0, bytes.to_vec()) + }; + let reference = PayloadReference { + offset, + stored: stored.len() as u32, + decoded: bytes.len() as u32, + checksum: codec::checksum(&stored), + encoding, + compression, + }; + reference.validate(offset, reference.end()?, page)?; + Ok((reference, stored)) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct Bounds { + pub first: MetricSeriesId, + pub last: MetricSeriesId, + pub min: i64, + pub max: i64, +} +impl Bounds { + pub fn validate(&self) -> Result<()> { + require( + self.first <= self.last && self.min <= self.max, + "reversed sample bounds", + ) + } + pub fn extend(&mut self, next: Self) { + self.last = next.last; + self.min = self.min.min(next.min); + self.max = self.max.max(next.max); + } + fn encode(&self, out: &mut Vec) { + for id in [self.first, self.last] { + put_u32(out, id.table_id); + put_u64(out, id.tsid); + } + put_u64(out, self.min as u64); + put_u64(out, self.max as u64); + } + fn decode(c: &mut Cursor<'_>) -> Result { + let result = Self { + first: MetricSeriesId { + table_id: c.u32()?, + tsid: c.u64()?, + }, + last: MetricSeriesId { + table_id: c.u32()?, + tsid: c.u64()?, + }, + min: c.u64()? as i64, + max: c.u64()? as i64, + }; + result.validate()?; + Ok(result) + } +} +#[derive(Debug, Clone)] +pub(super) struct PageReference { + pub bounds: Bounds, + pub payload: PayloadReference, +} +impl PageReference { + fn encode(&self, out: &mut Vec) { + self.bounds.encode(out); + self.payload.encode(out); + } + fn decode(c: &mut Cursor<'_>) -> Result { + let reference = Self { + bounds: Bounds::decode(c)?, + payload: PayloadReference::decode(c)?, + }; + if reference.payload.encoding != 1 || reference.payload.compression > 1 { + return Err(unsupported("samples metadata codec")); + } + Ok(reference) + } +} +#[derive(Debug, Clone)] +pub(super) struct Column { + pub role: u32, + pub id: u32, + pub payload: PayloadReference, + pub nulls: u32, +} +#[derive(Debug, Clone)] +pub(super) struct Block { + pub id: MetricSeriesId, + pub min: i64, + pub max: i64, + pub ordinal: u64, + pub rows: u32, + pub columns: Vec, +} +impl Block { + pub fn bounds(&self) -> Bounds { + Bounds { + first: self.id, + last: self.id, + min: self.min, + 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); + put_u64(out, self.min as u64); + put_u64(out, self.max as u64); + put_u64(out, self.ordinal); + put_u32(out, self.rows); + put_u32(out, self.columns.len() as u32); + for column in &self.columns { + put_u32(out, column.role); + put_u32(out, column.id); + column.payload.encode(out); + put_u32(out, column.nulls); + } + } + fn decode( + c: &mut Cursor<'_>, + manifest: &Manifest, + payload_end: u64, + previous_end: &mut u64, + ) -> Result { + let id = MetricSeriesId { + table_id: c.u32()?, + tsid: c.u64()?, + }; + let min = c.u64()? as i64; + let max = c.u64()? as i64; + let ordinal = c.u64()?; + let rows = c.u32()?; + let count = c.u32()? as usize; + require( + min <= max && valid_rows(rows as usize, manifest.fields.len()), + "invalid sample block bounds/count", + )?; + require( + count >= 3 && count <= manifest.fields.len() + 3 && count <= c.remaining() / 44, + "invalid sample column count", + )?; + let mut columns = Vec::with_capacity(count); + let mut next_field = 0; + for index in 0..count { + let role = c.u32()?; + let column_id = c.u32()?; + let payload = PayloadReference::decode(c)?; + let nulls = c.u32()?; + if index < 3 { + require( + role == index as u32 + 1 + && column_id + == if index == 0 { + manifest.timestamp.column_id + } else { + 0 + } + && nulls == 0, + "invalid mandatory sample column", + )?; + } else { + require( + role == 4 && nulls < rows, + "invalid sample field role/null count", + )?; + let position = manifest.fields[next_field..] + .iter() + .position(|f| f.column_id == column_id) + .ok_or_else(|| invalid("unknown or unordered sample field"))? + + next_field; + next_field = position + 1; + } + if payload.encoding <= 7 { + let valid = match role { + 1 => matches!(payload.encoding, 1 | 2), + 2 => matches!(payload.encoding, 3 | 4), + 3 => matches!(payload.encoding, 6 | 7), + 4 => payload.encoding == 5, + _ => false, + }; + require(valid, "sample codec role mismatch")?; + } + payload.validate(*previous_end, payload_end, false)?; + *previous_end = payload.end()?; + columns.push(Column { + role, + id: column_id, + payload, + nulls, + }); + } + Ok(Self { + id, + min, + max, + ordinal, + rows, + columns, + }) + } + pub fn follows(&self, previous: &Self, 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 { + Leaf(Vec), + Internal(Vec), +} +impl Page { + pub fn bounds(&self) -> Result { + let bounds: Vec<_> = match self { + Self::Leaf(v) => v.iter().map(Block::bounds).collect(), + Self::Internal(v) => v.iter().map(|c| c.bounds).collect(), + }; + let mut result = *bounds.first().ok_or_else(|| invalid("empty sample page"))?; + for bound in &bounds[1..] { + result.extend(*bound); + } + Ok(result) + } + pub fn encode(&self) -> Vec { + let mut out = Vec::new(); + match self { + Self::Leaf(v) => { + put_u32(&mut out, 0); + put_u64(&mut out, v.len() as u64); + for block in v { + block.encode(&mut out); + } + } + Self::Internal(v) => { + put_u32(&mut out, 1); + put_u64(&mut out, v.len() as u64); + for child in v { + child.encode(&mut out); + } + } + } + out + } + pub fn decode( + bytes: &[u8], + reference: &PageReference, + depth: u32, + payload_end: u64, + manifest: &Manifest, + ) -> Result { + require(bytes.len() <= PAGE_LIMIT, "sample page exceeds limit")?; + let mut c = Cursor::new(bytes); + let kind = c.u32()?; + let count = c.length()?; + require(count > 0, "empty sample page")?; + let page = match (kind, depth) { + (0, 1) => { + require( + count <= c.remaining() / 176, + "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)?; + if let Some(previous) = blocks.last() { + block.follows(previous, true)?; + } + blocks.push(block); + } + Self::Leaf(blocks) + } + (1, 2) => { + require( + count <= ROOT_CAPACITY && count <= c.remaining() / 72, + "sample root exceeds capacity", + )?; + let mut children: Vec = Vec::with_capacity(count); + let mut end = payload_end; + for _ in 0..count { + let child = PageReference::decode(&mut c)?; + child + .payload + .validate(end, reference.payload.offset, true)?; + end = child.payload.end()?; + if let Some(previous) = children.last() { + require( + previous.bounds.last <= child.bounds.first, + "unordered sample child ranges", + )?; + } + children.push(child); + } + Self::Internal(children) + } + _ => return Err(unsupported("samples metadata kind/depth")), + }; + c.end()?; + require( + page.bounds()? == reference.bounds, + "samples page bounds disagree with parent", + )?; + Ok(page) + } +} +pub(super) fn header() -> Vec { + let mut out = MAGIC.to_vec(); + put_u32(&mut out, VERSION); + put_u32(&mut out, 0); + out +} +pub(super) fn validate_header(bytes: &[u8]) -> Result<()> { + let mut c = Cursor::new(bytes); + require(c.take(8)? == MAGIC, "sample header magic")?; + if c.u32()? != VERSION { + return Err(unsupported("samples file version")); + } + require(c.u32()? == 0, "nonzero sample header reserved field")?; + c.end() +} +pub(super) struct Trailer { + pub root: PageReference, + pub depth: u32, + pub file_length: u64, + pub payload_end: u64, +} +impl Trailer { + pub fn encode(&self) -> Vec { + let mut out = Vec::new(); + self.root.encode(&mut out); + put_u32(&mut out, self.depth); + put_u32(&mut out, VERSION); + put_u64(&mut out, self.file_length); + put_u64(&mut out, self.payload_end); + out.extend(MAGIC); + let sum = codec::checksum(&out); + put_u64(&mut out, sum); + out + } + pub fn decode(bytes: &[u8], file_length: u64) -> Result { + require( + bytes.len() == TRAILER_SIZE as usize, + "truncated sample trailer", + )?; + let mut c = Cursor::new(bytes); + let root = PageReference::decode(&mut c)?; + let depth = c.u32()?; + let version = c.u32()?; + if !(1..=2).contains(&depth) || version != VERSION { + return Err(unsupported("sample trailer version/depth")); + } + let length = c.u64()?; + let payload_end = c.u64()?; + require( + c.take(8)? == MAGIC && c.u64()? == codec::checksum(&bytes[..104]), + "sample trailer checksum/magic", + )?; + require( + length == file_length + && payload_end >= HEADER_SIZE + && file_length >= HEADER_SIZE + TRAILER_SIZE, + "invalid sample framing", + )?; + root.payload + .validate(payload_end, file_length - TRAILER_SIZE, true)?; + require( + root.payload.end()? == file_length - TRAILER_SIZE, + "sample root must end at trailer", + )?; + Ok(Self { + root, + depth, + file_length, + payload_end, + }) + } +} diff --git a/src/mito2/src/series_index/samples/reader.rs b/src/mito2/src/series_index/samples/reader.rs new file mode 100644 index 00000000000..36e74d406d9 --- /dev/null +++ b/src/mito2/src/series_index/samples/reader.rs @@ -0,0 +1,418 @@ +// 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. + +//! Lazy, selective sample reads with independent traversal and pinned lifetimes. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + +use datatypes::arrow::array::{ + ArrayRef, Float64Array, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt32Array, UInt64Array, +}; +use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use futures::stream::BoxStream; +use object_store::ObjectStore; +use snafu::ResultExt; + +use crate::error::{NewRecordBatchSnafu, OpenDalSnafu, Result, SerdeJsonSnafu}; +use crate::series_index::artifact::{ArtifactFormat, SamplesFileHandle}; +use crate::series_index::format::{SamplesFormatReader, SamplesReadRequest}; +use crate::series_index::samples::codec; +use crate::series_index::samples::metadata::{ + self, Block, Bounds, Column, Page, PageReference, Trailer, +}; +use crate::series_index::samples::{ + MANIFEST_LIMIT, Manifest, MetricSeriesId, TimestampUnit, VERSION, invalid, require, unsupported, +}; +use crate::sst::parquet::helper::fetch_byte_ranges; + +pub(super) struct SamplesReader { + store: ObjectStore, + manifest: Arc, + handle: SamplesFileHandle, +} +impl SamplesReader { + pub async fn open(store: ObjectStore, handle: SamplesFileHandle) -> Result { + let entry = handle.entry(); + let artifact = &entry.artifact; + require( + artifact.format == ArtifactFormat::ColumnarSamples, + "samples artifact format mismatch", + )?; + let region = handle.file_id().region_id().as_u64(); + require( + handle.file_id().file_id() == entry.file_id, + "samples source handle mismatch", + )?; + let path = format!( + "{region}/samples/{}/{}/manifest.json", + entry.file_id, artifact.generation + ); + let mut objects = HashMap::new(); + for object in &artifact.objects { + require( + objects + .insert(object.path.as_str(), object.file_size) + .is_none(), + "duplicate sample artifact object", + )?; + } + let length = *objects + .get(path.as_str()) + .ok_or_else(|| invalid("missing samples manifest"))?; + require( + length > 0 && length <= MANIFEST_LIMIT as u64, + "samples manifest exceeds limit", + )?; + let bytes = read_range(&store, &path, 0, length).await?; + let manifest: Manifest = serde_json::from_slice(&bytes).context(SerdeJsonSnafu)?; + if manifest.version != VERSION { + return Err(unsupported("samples manifest version")); + } + require( + manifest.format == artifact.format + && manifest.region_id == region + && manifest.source_file_id == entry.file_id + && manifest.generation == artifact.generation, + "samples manifest identity mismatch", + )?; + manifest.validate_schema()?; + let mut paths = HashSet::from([path]); + let mut total = length; + let mut previous = None; + for (index, part) in manifest.files.iter().enumerate() { + require( + index as u64 == u64::from(part.part_id) + && part.num_rows > 0 + && part.num_blocks > 0 + && part.num_blocks <= part.num_rows + && part.first <= part.last + && part.min_timestamp <= part.max_timestamp + && part.file_size >= metadata::HEADER_SIZE + metadata::TRAILER_SIZE, + "invalid sample part descriptor", + )?; + if let Some(last) = previous { + require(last <= part.first, "unordered sample parts")?; + } + previous = Some(part.last); + let part_path = manifest.part_path(part.part_id); + require( + objects.get(part_path.as_str()) == Some(&part.file_size), + "sample part size mismatch", + )?; + paths.insert(part_path); + total = total + .checked_add(part.file_size) + .ok_or_else(|| invalid("sample artifact size overflow"))?; + } + require( + total == artifact.file_size + && paths.len() == artifact.objects.len() + && artifact.objects.iter().all(|o| paths.contains(&o.path)), + "sample artifact objects/size mismatch", + )?; + Ok(Self { + store, + manifest: Arc::new(manifest), + handle, + }) + } +} +impl SamplesFormatReader for SamplesReader { + fn schema(&self) -> SchemaRef { + self.manifest.schema() + } + fn read(&self, request: SamplesReadRequest) -> Result>> { + let schema = self.schema(); + let mut projection = Vec::new(); + let mut seen = HashSet::new(); + for name in &request.projection { + require(seen.insert(name), "duplicate sample projection")?; + projection.push( + schema + .index_of(name) + .map_err(|_| invalid(format!("unknown sample projection {name}")))?, + ); + } + require( + request + .start + .zip(request.end) + .is_none_or(|(start, end)| start <= end), + "reversed sample time range", + )?; + let empty = request.ids.as_ref().is_some_and(|ids| ids.is_empty()) + || request.start.zip(request.end).is_some_and(|(a, b)| a == b); + let complete = request.ids.is_none() + && request.table_id.is_none() + && request.start.is_none() + && request.end.is_none(); + let projected = Arc::new( + schema + .project(&projection) + .map_err(|e| invalid(format!("sample projection: {e}")))?, + ); + let manifest = self.manifest.clone(); + let store = self.store.clone(); + let handle = self.handle.clone(); + Ok(Box::pin(async_stream::try_stream! { + let _handle = handle; + if empty { return; } + let mut previous: Option = None; + let mut adjacent = true; + 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; } + let path = manifest.part_path(part.part_id); + let header = read_range(&store, &path, 0, metadata::HEADER_SIZE).await?; + metadata::validate_header(&header)?; + let bytes = read_range(&store, &path, part.file_size - metadata::TRAILER_SIZE, part.file_size).await?; + let trailer = Trailer::decode(&bytes, part.file_size)?; + require(trailer.root.bounds == bounds, "sample root/part bounds mismatch")?; + let root = load_metadata_page(&store, &path, &trailer.root, trailer.depth, trailer.payload_end, &manifest).await?; + let (mut active, children) = match root { Page::Leaf(blocks) => (Some(blocks), Vec::new()), Page::Internal(children) => (None, children) }; + let mut children = children.into_iter(); let mut row_count = 0u64; let mut block_count = 0u64; + let mut payload_position = metadata::HEADER_SIZE; + loop { + let blocks = if let Some(blocks) = active.take() { blocks } else { + let Some(child) = children.next() else { break; }; + if !matches(child.bounds, &request) { adjacent = false; continue; } + match load_metadata_page(&store, &path, &child, 1, trailer.payload_end, &manifest).await? { + Page::Leaf(blocks) => blocks, + Page::Internal(_) => Err(unsupported("internal child in sample depth-two tree"))?, + } + }; + for block in blocks { + 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")?; } + 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) { + read_block(&store, &path, &block, &manifest, &request, &projection, projected.clone()).await? + } else { None }; + if let Some(batch) = batch { yield batch; } + previous = Some(block); adjacent = true; + } + } + if complete { require(row_count == part.num_rows && block_count == part.num_blocks, "sample part totals mismatch")?; } + } + })) + } +} +fn matches(bounds: Bounds, request: &SamplesReadRequest) -> bool { + if request.start.is_some_and(|start| bounds.max < start) + || request.end.is_some_and(|end| bounds.min >= end) + { + return false; + } + let mut first = bounds.first; + let mut last = bounds.last; + if let Some(table_id) = request.table_id { + first = first.max(MetricSeriesId { table_id, tsid: 0 }); + last = last.min(MetricSeriesId { + table_id, + tsid: u64::MAX, + }); + } + first <= last + && request + .ids + .as_ref() + .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")?; + let mut buffers = fetch_byte_ranges(path, store.clone(), std::slice::from_ref(&(start..end))) + .await + .context(OpenDalSnafu)?; + let bytes = buffers + .pop() + .ok_or_else(|| invalid("missing sample range response"))?; + require( + buffers.is_empty() && bytes.len() as u64 == end - start, + "sample range length mismatch", + )?; + Ok(bytes) +} +/// Cache integration can key this entrypoint by immutable generation, range and codecs. +async fn load_metadata_page( + store: &ObjectStore, + path: &str, + reference: &PageReference, + depth: u32, + payload_end: u64, + manifest: &Manifest, +) -> Result { + if reference.payload.encoding != 1 { + return Err(unsupported("sample metadata codec")); + } + let bytes = read_range( + store, + path, + reference.payload.offset, + reference.payload.end()?, + ) + .await?; + Page::decode( + &reference.payload.unpack(&bytes)?, + reference, + depth, + payload_end, + manifest, + ) +} +/// Keep column loads separate from metadata for future Mito cache integration. +async fn load_column_chunk(store: &ObjectStore, path: &str, column: &Column) -> Result> { + let allowed = match column.role { + 1 => matches!(column.payload.encoding, 1 | 2), + 2 => matches!(column.payload.encoding, 3 | 4), + 3 => matches!(column.payload.encoding, 6 | 7), + 4 => column.payload.encoding == 5, + _ => false, + }; + if !allowed { + return Err(unsupported("samples column encoding")); + } + let bytes = read_range(store, path, column.payload.offset, column.payload.end()?).await?; + column.payload.unpack(&bytes) +} +fn timestamp_array(values: impl Iterator, unit: TimestampUnit) -> ArrayRef { + match unit { + TimestampUnit::Second => Arc::new(TimestampSecondArray::from_iter_values(values)), + TimestampUnit::Millisecond => Arc::new(TimestampMillisecondArray::from_iter_values(values)), + TimestampUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from_iter_values(values)), + TimestampUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from_iter_values(values)), + } +} +async fn read_block( + store: &ObjectStore, + path: &str, + block: &Block, + manifest: &Manifest, + request: &SamplesReadRequest, + projection: &[usize], + schema: SchemaRef, +) -> Result> { + let mut selected: Vec = (0..block.rows as usize).collect(); + let mut times = None; + if request.start.is_some() || request.end.is_some() || projection.contains(&2) { + let column = &block.columns[0]; + let values = codec::decode_numbers( + &load_column_chunk(store, path, column).await?, + column.payload.encoding, + block.rows as usize, + )?; + require( + values.first().map(|v| *v as i64) == Some(block.min) + && values.last().map(|v| *v as i64) == Some(block.max), + "sample timestamp bounds mismatch", + )?; + selected.retain(|&i| { + request + .start + .is_none_or(|start| (values[i] as i64) >= start) + && request.end.is_none_or(|end| (values[i] as i64) < end) + }); + if selected.is_empty() { + return Ok(None); + } + times = Some(values); + } + let mut arrays: BTreeMap = BTreeMap::new(); + // Visit columns in physical order, regardless of requested output order. + let mut sorted = projection.to_vec(); + sorted.sort_unstable(); + 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(store, path, column).await?, + 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", + )?; + } + 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(store, path, column).await?, + 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( + &load_column_chunk(store, path, column).await?, + 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()])) + } + } + }; + arrays.insert(index, array); + } + let output = projection + .iter() + .map(|i| { + arrays + .remove(i) + .ok_or_else(|| invalid("missing sample projected column")) + }) + .collect::>>()?; + Ok(Some( + RecordBatch::try_new_with_options( + schema, + output, + &RecordBatchOptions::new().with_row_count(Some(selected.len())), + ) + .context(NewRecordBatchSnafu)?, + )) +} diff --git a/src/mito2/src/series_index/samples/tests.rs b/src/mito2/src/series_index/samples/tests.rs new file mode 100644 index 00000000000..8048d1aebfb --- /dev/null +++ b/src/mito2/src/series_index/samples/tests.rs @@ -0,0 +1,1065 @@ +// 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. + +//! Lossless codecs, selective time navigation, corruption and artifact lifecycle. + +use std::collections::BTreeSet; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use datatypes::arrow::array::{ + Array, ArrayRef, BinaryArray, BinaryDictionaryBuilder, Float64Array, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use futures::TryStreamExt; +use object_store::ObjectStore; +use object_store::layers::mock::{self, MockLayerBuilder, oio}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::FileId; + +use super::metadata::{Bounds, Page, Trailer}; +use super::*; +use crate::series_index::artifact::{SamplesEntry, SamplesFileHandle}; +use crate::series_index::format::{SamplesFormat, SamplesReadRequest}; +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}; + +type Sample = (u32, u64, i64, u64, u8, Option); +fn metadata() -> RegionMetadataRef { + let mut m = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + m.column_metadatas + .iter_mut() + .find(|c| c.column_id == 2) + .unwrap() + .column_schema + .data_type = ConcreteDataType::float64_datatype(); + Arc::new(m) +} +fn store() -> ObjectStore { + ObjectStore::new(object_store::services::Memory::default()).unwrap() +} +fn batch(meta: &RegionMetadataRef, rows: &[Sample], dictionary: bool) -> RecordBatch { + let keys: Vec<_> = rows + .iter() + .map(|r| new_sparse_primary_key(&["a", "b"], meta, r.0, r.1)) + .collect(); + let keys: ArrayRef = if dictionary { + let mut builder = BinaryDictionaryBuilder::::new(); + for key in &keys { + builder.append(key).unwrap(); + } + Arc::new(builder.finish()) + } else { + Arc::new(BinaryArray::from_iter_values( + keys.iter().map(Vec::as_slice), + )) + }; + let unit = meta + .time_index_column() + .column_schema + .data_type + .as_timestamp() + .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), + ), + ), + }; + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("field_0", DataType::Float64, true), + Field::new("ts", ts.data_type().clone(), false), + Field::new("__primary_key", keys.data_type().clone(), false), + Field::new(SEQUENCE, DataType::UInt64, false), + Field::new(OPERATION, DataType::UInt8, false), + ])), + vec![ + Arc::new(Float64Array::from_iter( + rows.iter().map(|r| r.5.map(f64::from_bits)), + )), + 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))), + ], + ) + .unwrap() +} +async fn build( + format: &ColumnarSamplesFormat, + store: &ObjectStore, + meta: RegionMetadataRef, + batches: &[RecordBatch], +) -> SamplesFileHandle { + let file_id = FileId::random(); + let generation = FileId::random(); + let prefix = format!( + "{}/samples/{file_id}/{generation}/", + meta.region_id.as_u64() + ); + let mut writer = format + .writer(meta.clone(), store.clone(), &prefix, file_id, generation) + .await + .unwrap(); + for batch in batches { + writer.write(batch).await.unwrap(); + } + let artifact = writer.finish().await.unwrap(); + let (purger, _) = series_index_channel(store.clone()); + SamplesFileHandle::new( + RegionFileId::new(meta.region_id, file_id), + SamplesEntry { file_id, artifact }, + purger, + ) +} +async fn manifest(store: &ObjectStore, handle: &SamplesFileHandle) -> Manifest { + let path = &handle + .entry() + .artifact + .objects + .iter() + .find(|o| o.path.ends_with("manifest.json")) + .unwrap() + .path; + serde_json::from_slice(&store.read(path).await.unwrap().to_bytes()).unwrap() +} +fn request(names: &[&str]) -> SamplesReadRequest { + SamplesReadRequest { + projection: names.iter().map(|n| n.to_string()).collect(), + ..Default::default() + } +} +fn full_request() -> SamplesReadRequest { + request(&[ + TABLE_ID_COLUMN, + TSID_COLUMN, + "ts", + SEQUENCE, + OPERATION, + "field_0", + ]) +} +fn collect_rows(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| { + let ids = b + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let tsids = b.column(1).as_any().downcast_ref::().unwrap(); + let (times, _) = + datatypes::timestamp::timestamp_array_to_primitive(b.column(2)).unwrap(); + let seq = b.column(3).as_any().downcast_ref::().unwrap(); + let op = b.column(4).as_any().downcast_ref::().unwrap(); + let values = b.column(5).as_any().downcast_ref::().unwrap(); + (0..b.num_rows()) + .map(|i| { + ( + ids.value(i), + tsids.value(i), + times.value(i), + seq.value(i), + op.value(i), + if values.is_null(i) { + None + } else { + Some(values.value(i).to_bits()) + }, + ) + }) + .collect::>() + }) + .collect() +} + +#[test] +fn sample_codecs_preserve_bits_and_reject_corruption() { + for times in [ + vec![0], + vec![1, 1, 1], + vec![-10, -5, 0], + vec![i64::MIN, i64::MAX], + vec![-100, 0, 1, 999], + ] { + let (id, bytes) = codec::encode_time(×); + assert_eq!( + codec::decode_numbers(&bytes, id, times.len()) + .unwrap() + .iter() + .map(|v| *v as i64) + .collect::>(), + times + ); + assert!(codec::decode_numbers(&bytes[..bytes.len() - 1], id, times.len()).is_err()); + } + for seq in [vec![0, 0, 0], vec![u64::MAX, 0, 1 << 63, 3, 2]] { + let (id, bytes) = codec::encode_sequence(&seq); + assert_eq!(codec::decode_numbers(&bytes, id, seq.len()).unwrap(), seq); + } + assert_eq!( + codec::encode_sequence(&[10, 9, 7]), + (4, vec![10, 0, 0, 0, 0, 0, 0, 0, 1, 3]) + ); + let bits = vec![ + Some(0), + None, + Some(1 << 63), + Some(0x7ff8000000000042), + None, + Some(u64::MAX), + Some(1), + ]; + let bytes = codec::encode_float(&bits); + assert_eq!(codec::decode_float(&bytes, bits.len(), 2).unwrap(), bits); + let mut bad = bytes.clone(); + bad[1] |= 128; + assert!(codec::decode_float(&bad, bits.len(), 2).is_err()); + assert!(codec::decode_float(&bytes, bits.len(), 1).is_err()); + let mut bad = bytes.clone(); + bad.push(0); + assert!(codec::decode_float(&bad, bits.len(), 2).is_err()); + assert!(codec::decode_numbers(&[255; 18], 4, 2).is_err()); + assert!(codec::decode_operations(&[2], 6, 2).is_err()); + assert!(codec::decode_operations(&[1], 7, 2).is_err()); + assert!(codec::decode_float(&[0], 8193, 0).is_err()); + for ops in [vec![1, 1], vec![0, 1, 0]] { + let (id, bytes) = codec::encode_operations(&ops); + assert_eq!( + codec::decode_operations(&bytes, id, ops.len()).unwrap(), + ops + ); + } +} +#[tokio::test] +async fn lossless_round_trip_all_units_and_input_batches() { + for timestamp_type in [ + ConcreteDataType::timestamp_second_datatype(), + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_microsecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ] { + let mut m = (*metadata()).clone(); + m.column_metadatas + .iter_mut() + .find(|c| c.column_id == 3) + .unwrap() + .column_schema + .data_type = timestamp_type; + let m = Arc::new(m); + let store = store(); + let format = ColumnarSamplesFormat { + block_rows: 2, + ..Default::default() + }; + let rows = vec![ + (1, 42, -1, 9, 1, Some(1 << 63)), + (1, 42, 0, 8, 0, None), + (1, 42, 0, 8, 1, Some(0x7ff8000000000042)), + (1, 42, 0, 7, 1, Some(u64::MAX)), + (2, 42, 5, u64::MAX, 1, None), + ]; + let handle = build( + &format, + &store, + m.clone(), + &[batch(&m, &rows[..2], false), batch(&m, &rows[2..], true)], + ) + .await; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let result = reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_rows(&result), rows); + assert!(result.iter().all(|b| b.num_rows() <= 2)); + let no_columns = reader + .read(request(&[])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + no_columns.iter().map(|b| b.num_rows()).sum::(), + rows.len() + ); + for object in &handle.entry().artifact.objects { + assert_eq!( + store.read(&object.path).await.unwrap().len() as u64, + object.file_size + ); + } + } +} +#[tokio::test] +async fn equal_timestamps_cross_leaves_and_parts() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat { + block_rows: 1, + page_bytes: 232, + root_capacity: 2, + ..Default::default() + }; + let rows: Vec<_> = (0..7) + .map(|i| (7, 42, 100, 100 - i, 1, Some(i))) + .chain([(7, 43, 50, 1, 1, Some(9)), (8, 42, 100, 1, 1, None)]) + .collect(); + let handle = build(&format, &store, m.clone(), &[batch(&m, &rows, true)]).await; + let manifest = manifest(&store, &handle).await; + 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, + 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.start = Some(start); + req.end = Some(end); + let out = reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + collect_rows(&out), + rows.iter() + .copied() + .filter(|r| r.0 == 7 && r.1 == 42 && r.2 >= start && r.2 < end) + .collect::>() + ); + } + let mut req = request(&["field_0", TSID_COLUMN]); + req.table_id = Some(8); + let out = reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(out.len(), 1); + assert_eq!(out[0].column(0).null_count(), 1); +} +#[tokio::test] +async fn empty_selection_and_validation() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let handle = build(&format, &store, m, &[]).await; + assert!(manifest(&store, &handle).await.files.is_empty()); + assert_eq!(handle.entry().artifact.objects.len(), 1); + let reader = format.reader(store, handle).await.unwrap(); + assert!( + reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + assert!(reader.read(request(&["unknown"])).is_err()); + assert!(reader.read(request(&["ts", "ts"])).is_err()); + let mut req = request(&[]); + req.start = Some(2); + req.end = Some(1); + assert!(reader.read(req).is_err()); + let mut req = request(&[]); + req.ids = Some(BTreeSet::new()); + assert!( + reader + .read(req) + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); +} +#[tokio::test] +async fn failed_writer_and_abort_remove_outputs() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat { + block_rows: 1, + part_bytes: 1, + ..Default::default() + }; + for drop_writer in [false, true] { + let id = FileId::random(); + let generation = FileId::random(); + let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + let mut writer = format + .writer(m.clone(), store.clone(), &prefix, id, generation) + .await + .unwrap(); + writer + .write(&batch( + &m, + &[(1, 1, 1, 1, 1, None), (1, 1, 2, 1, 1, Some(1))], + false, + )) + .await + .unwrap(); + if drop_writer { + drop(writer); + } else { + assert!( + writer + .write(&batch(&m, &[(1, 1, 0, 1, 1, None)], false)) + .await + .is_err() + ); + assert!(writer.finish().await.is_err()); + } + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if store + .list_with(&prefix) + .recursive(true) + .await + .unwrap() + .iter() + .all(|e| e.metadata().is_dir()) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + assert!( + format + .writer(m, store, "wrong", FileId::random(), FileId::random()) + .await + .is_err() + ); +} + +struct RecordingRead { + inner: oio::Reader, + path: String, + reads: Arc>>, +} +impl mock::Read for RecordingRead { + async fn open( + &self, + range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, Box)> { + self.reads.lock().unwrap().push((self.path.clone(), range)); + self.inner.open(range).await + } + async fn read(&self, range: mock::BytesRange) -> mock::Result<(mock::RpRead, mock::Buffer)> { + self.reads.lock().unwrap().push((self.path.clone(), range)); + self.inner.read(range).await + } +} +#[tokio::test] +async fn selective_reads_pin_the_artifact() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let built = build( + &format, + &store, + m.clone(), + &[batch( + &m, + &[(1, 1, 0, 1, 1, Some(7)), (2, 1, 10, 1, 1, None)], + false, + )], + ) + .await; + let (purger, mut purges) = series_index_channel(store.clone()); + let handle = SamplesFileHandle::new(built.file_id(), built.entry().clone(), purger); + drop(built); + let reads = Arc::new(Mutex::new(Vec::new())); + let observed = reads.clone(); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| { + Box::new(RecordingRead { + inner, + path: path.into(), + reads: observed.clone(), + }) + })) + .build() + .unwrap(); + let reader = format + .reader(store.layer(layer), handle.clone()) + .await + .unwrap(); + reads.lock().unwrap().clear(); + let mut req = request(&["field_0"]); + req.table_id = Some(2); + let stream = reader.read(req).unwrap(); + handle.mark_deleted(); + drop(handle); + drop(reader); + assert!(purges.try_recv().is_err()); + let result = stream.try_collect::>().await.unwrap(); + 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); +} +#[test] +fn payload_limits_and_timestamp_envelopes() { + assert!(valid_rows(8192, 1)); + assert!(!valid_rows(8193, 1)); + assert!(valid_rows(31, 1024)); + assert!(!valid_rows(32, 1024)); + let mut bounds = Bounds { + first: MetricSeriesId { + table_id: 1, + tsid: 1, + }, + last: MetricSeriesId { + table_id: 1, + tsid: 1, + }, + min: 100, + max: 200, + }; + bounds.extend(Bounds { + first: MetricSeriesId { + table_id: 1, + tsid: 2, + }, + last: MetricSeriesId { + table_id: 1, + tsid: 2, + }, + min: 0, + max: 50, + }); + assert_eq!((bounds.min, bounds.max), (0, 200)); + let (mut reference, bytes) = metadata::pack(&[1, 2, 3], (1u64 << 32) + 5, 1, false).unwrap(); + reference + .validate(16, reference.end().unwrap(), false) + .unwrap(); + assert_eq!(reference.unpack(&bytes).unwrap(), [1, 2, 3]); + reference.offset = u64::MAX; + assert!(reference.end().is_err()); + let (mut reference, bytes) = metadata::pack(&vec![0; 1024], 16, 1, false).unwrap(); + reference.decoded = 1000; + assert!(reference.unpack(&bytes).is_err()); +} +#[tokio::test] +async fn malformed_framing_and_payload_corruption() { + let store = store(); + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + let handle = build( + &format, + &store, + m.clone(), + &[batch(&m, &[(1, 1, 0, 1, 1, Some(42))], false)], + ) + .await; + 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 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(); + unsupported.root.payload.encoding = encoding; + unsupported.root.payload.compression = compression; + assert!(Trailer::decode(&unsupported.encode(), original.len() as u64).is_err()); + } + let decoded = trailer + .root + .payload + .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() + else { + panic!() + }; + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let mut bad = original.clone(); + bad[blocks[0].columns[3].payload.offset as usize] ^= 1; + store.write(&path, bad).await.unwrap(); + assert!( + reader + .read(request(&[TABLE_ID_COLUMN])) + .unwrap() + .try_collect::>() + .await + .is_ok() + ); + assert!( + reader + .read(request(&["field_0"])) + .unwrap() + .try_collect::>() + .await + .is_err() + ); + for offset in [0, end + 72, end + 104] { + let mut bad = original.clone(); + bad[offset] ^= 3; + store.write(&path, bad).await.unwrap(); + assert!( + reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .is_err() + ); + } + for length in 0..decoded.len() { + assert!( + Page::decode( + &decoded[..length], + &trailer.root, + 1, + trailer.payload_end, + &manifest + ) + .is_err() + ); + } + let mut entry = handle.entry().clone(); + entry + .artifact + .objects + .push(entry.artifact.objects[0].clone()); + let (purger, _) = series_index_channel(store.clone()); + let bad = SamplesFileHandle::new(handle.file_id(), entry, purger); + assert!(format.reader(store, bad).await.is_err()); +} + +#[tokio::test] +async fn real_samples_reconcile_reopen_and_retire() { + use crate::config::MitoConfig; + use crate::series_index::format::SeriesFormats; + use crate::series_index::maintenance::reconcile_series_indexes; + use crate::test_util::TestEnv; + let mut env = TestEnv::with_prefix("columnar-samples-reconcile").await; + let (engine, region) = crate::series_index::tests::prepare_region_with_metadata( + &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(), + ) + .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, mut 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!(reconcile().await.unwrap().built_samples, 2); + assert_eq!(reconcile().await.unwrap().built_samples, 0); + let reopened = + crate::series_index::catalog::load_version_control(&store, region.region_id, &purger).await; + assert_eq!(reopened.current().samples.len(), 2); + let mut timestamps = Vec::new(); + for handle in reopened.current().samples.values() { + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let batches = reader + .read(full_request()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + timestamps.extend(collect_rows(&batches).into_iter().map(|row| row.2)); + } + timestamps.sort_unstable(); + assert_eq!(timestamps, vec![1000, 2000]); + // A missing artifact triggers rebuilding; failure to publish its replacement + // must retain the prior snapshot/catalog and retire only the new generation. + let before = region.series_index_version(); + let catalog_path = crate::series_index::catalog::samples_catalog_path(region.region_id); + let catalog = store.read(&catalog_path).await.unwrap().to_bytes(); + let old = before.samples.values().next().unwrap(); + store + .delete(&old.entry().artifact.objects[0].path) + .await + .unwrap(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(|path, _, inner| { + if path.ends_with("samples.json") { + Box::new(FaultyWriter { + inner, + fail_write: true, + closing: None, + }) + } else { + inner + } + })) + .build() + .unwrap(); + assert!( + reconcile_series_indexes( + 0, + store.clone().layer(layer), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true + ) + .await + .is_err() + ); + assert!(Arc::ptr_eq(&before, ®ion.series_index_version())); + assert_eq!(catalog, store.read(&catalog_path).await.unwrap().to_bytes()); + let cleanup = purges.try_recv().unwrap().objects.unwrap(); + assert_eq!(cleanup.len(), 2); + assert!(cleanup.iter().all(|path| { + !old.entry() + .artifact + .objects + .iter() + .any(|object| object.path == *path) + })); + assert!(purges.try_recv().is_err()); + use store_api::region_engine::RegionEngine; + engine.stop().await.unwrap(); +} + +async fn fixture(store: &ObjectStore) -> SamplesFileHandle { + let bytes = + include_str!("../../../../../tests/compatibility/wire/mito_columnar_samples/v1-part.hex") + .split_whitespace() + .collect::(); + let bytes = (0..bytes.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&bytes[i..i + 2], 16).unwrap()) + .collect::>(); + 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 objects = vec![ + crate::series_index::artifact::ArtifactObject { + path: manifest.path(), + file_size: json.len() as u64, + }, + crate::series_index::artifact::ArtifactObject { + path: manifest.part_path(0), + file_size: bytes.len() as u64, + }, + ]; + let artifact = crate::series_index::artifact::ArtifactMetadata { + format: ArtifactFormat::ColumnarSamples, + generation: manifest.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, + ), + SamplesEntry { + file_id: manifest.source_file_id, + artifact, + }, + purger, + ) +} +#[tokio::test] +async fn frozen_v1_samples_fixture_and_unknown_codec() { + let store = store(); + let handle = fixture(&store).await; + let format = ColumnarSamplesFormat::default(); + let reader = format.reader(store.clone(), handle.clone()).await.unwrap(); + let batches = reader + .read(full_request()) + .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!( + reader + .read(request(&["ts"])) + .unwrap() + .try_collect::>() + .await + .is_ok() + ); + assert!( + reader + .read(request(&["field_0"])) + .unwrap() + .try_collect::>() + .await + .is_err() + ); +} + +#[tokio::test] +async fn wide_schema_byte_bound_and_unsupported_input() { + use datatypes::schema::ColumnSchema; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + let m = metadata(); + let store = store(); + let format = ColumnarSamplesFormat::default(); + let mut builder = RegionMetadataBuilder::from_existing((*m).clone()); + for i in 1..1024 { + builder.push_column_metadata(ColumnMetadata { + column_id: 100 + i, + semantic_type: api::v1::SemanticType::Field, + column_schema: ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + ), + }); + } + let wide = Arc::new(builder.build().unwrap()); + let rows: Vec<_> = (0..65).map(|i| (1, 1, i, 1, 1, Some(i as u64))).collect(); + let input = batch(&wide, &rows, false); + let mut fields = input.schema().fields().to_vec(); + let mut arrays = input.columns().to_vec(); + for i in 1..1024 { + fields.insert( + i, + Arc::new(Field::new(format!("field_{i}"), DataType::Float64, true)), + ); + arrays.insert(i, input.column(0).clone()); + } + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap(); + let handle = build(&format, &store, wide, &[input]).await; + let batches = format + .reader(store.clone(), handle) + .await + .unwrap() + .read(request(&["field_1023", "ts"])) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + batches.iter().map(|b| b.num_rows()).collect::>(), + vec![31, 31, 3] + ); + let wrong = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let id = FileId::random(); + let generation = FileId::random(); + let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + assert!( + format + .writer(wrong, store.clone(), &prefix, id, generation) + .await + .is_err() + ); + let mut writer = format + .writer(m.clone(), store, &prefix, id, generation) + .await + .unwrap(); + let mut bad = (*m).clone(); + bad.column_metadatas + .iter_mut() + .find(|c| c.column_id == 3) + .unwrap() + .column_schema + .data_type = ConcreteDataType::timestamp_second_datatype(); + assert!( + writer + .write(&batch(&Arc::new(bad), &rows, false)) + .await + .is_err() + ); + writer.abort().await.unwrap(); +} + +struct FaultyWriter { + inner: oio::Writer, + fail_write: bool, + closing: Option>, +} +impl mock::Write for FaultyWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + if self.fail_write { + return Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected sample write failure", + )); + } + self.inner.write(buffer).await + } + async fn close(&mut self) -> mock::Result { + if let Some(closing) = &self.closing { + closing.notify_one(); + std::future::pending::<()>().await; + } + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected sample close failure", + )) + } + async fn abort(&mut self) -> mock::Result<()> { + self.inner.abort().await + } +} +#[tokio::test] +async fn output_failures_and_cancelled_finish_clean_up() { + let m = metadata(); + let format = ColumnarSamplesFormat::default(); + for (manifest_only, fail_write, cancel) in [ + (false, true, false), + (false, false, false), + (true, false, false), + (false, false, true), + ] { + let store = store(); + let closing = Arc::new(tokio::sync::Notify::new()); + let observed = closing.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if !manifest_only || path.ends_with("manifest.json") { + Box::new(FaultyWriter { + inner, + fail_write, + closing: cancel.then(|| observed.clone()), + }) + } else { + inner + } + })) + .build() + .unwrap(); + let id = FileId::random(); + let generation = FileId::random(); + let prefix = format!("{}/samples/{id}/{generation}/", m.region_id.as_u64()); + let mut writer = format + .writer( + m.clone(), + store.clone().layer(layer), + &prefix, + id, + generation, + ) + .await + .unwrap(); + writer + .write(&batch(&m, &[(1, 1, 0, 1, 1, Some(7))], false)) + .await + .unwrap(); + if cancel { + let task = tokio::spawn(async move { writer.finish().await }); + tokio::time::timeout(Duration::from_secs(5), closing.notified()) + .await + .unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + } else { + assert!(writer.finish().await.is_err()); + } + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if store + .list_with(&prefix) + .recursive(true) + .await + .unwrap() + .iter() + .all(|e| e.metadata().is_dir()) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } +} diff --git a/src/mito2/src/series_index/samples/writer.rs b/src/mito2/src/series_index/samples/writer.rs new file mode 100644 index 00000000000..5216ba8145e --- /dev/null +++ b/src/mito2/src/series_index/samples/writer.rs @@ -0,0 +1,485 @@ +// 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. + +//! Streaming per-SST sample writing with bounded blocks and owned outputs. + +use async_trait::async_trait; +use datatypes::arrow::array::{ + Array, BinaryArray, DictionaryArray, Float64Array, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::{DataType, UInt32Type}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::DataType as _; +use datatypes::timestamp::timestamp_array_to_primitive; +use mito_codec::row_converter::SparseOffsetsCache; +use object_store::ObjectStore; +use snafu::ResultExt; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::{FileId, consts}; + +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::samples::codec; +use crate::series_index::samples::metadata::{ + self, Block, Column, 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}; + +struct Row { + id: MetricSeriesId, + time: i64, + sequence: u64, + operation: u8, + fields: Vec>, +} +pub(super) struct SamplesWriter { + options: ColumnarSamplesFormat, + manifest: Manifest, + outputs: Outputs, + offsets: SparseOffsetsCache, + key_buffer: Vec, + rows: Vec, + last: Option<(MetricSeriesId, i64, u64)>, + ordinal: u64, + leaves: Vec>, + leaf_bytes: usize, + descriptor_bytes: usize, + position: u64, + part_rows: u64, + part_blocks: u64, + manifest_bytes: usize, + failed: bool, +} +impl SamplesWriter { + pub fn new( + options: ColumnarSamplesFormat, + metadata: RegionMetadataRef, + store: ObjectStore, + prefix: &str, + source_file_id: FileId, + generation: FileId, + ) -> Result { + validate_metadata(&metadata)?; + require( + (1..=super::MAX_ROWS).contains(&options.block_rows) + && (1..=metadata::ROOT_CAPACITY).contains(&options.root_capacity) + && (188..=metadata::PAGE_LIMIT).contains(&options.page_bytes) + && options.part_bytes > 0 + && options.metadata_bytes > 0, + "invalid samples packing options", + )?; + let timestamp = metadata.time_index_column(); + let DataType::Timestamp(unit, None) = timestamp.column_schema.data_type.as_arrow_type() + else { + return Err(invalid("samples require a timestamp time index")); + }; + let fields = metadata + .field_columns() + .map(|column| { + require( + column.column_schema.data_type.as_arrow_type() == DataType::Float64, + "samples support only Float64 fields", + )?; + Ok(ValueColumn { + column_id: column.column_id, + name: column.column_schema.name.clone(), + data_type: "float64".into(), + nullable: true, + }) + }) + .collect::>>()?; + let manifest = Manifest { + format: ArtifactFormat::ColumnarSamples, + version: VERSION, + region_id: metadata.region_id.as_u64(), + source_file_id, + generation, + timestamp: TimestampColumn { + column_id: timestamp.column_id, + name: timestamp.column_schema.name.clone(), + unit: TimestampUnit::from_arrow(&unit), + }, + fields, + files: Vec::new(), + }; + manifest.validate_schema()?; + require( + prefix == manifest.prefix(), + "samples output prefix mismatch", + )?; + let manifest_bytes = serde_json::to_vec(&manifest).context(SerdeJsonSnafu)?.len(); + require( + manifest_bytes <= MANIFEST_LIMIT, + "samples schema exceeds manifest limit", + )?; + Ok(Self { + options, + manifest, + outputs: Outputs { + store, + active: None, + paths: Vec::new(), + committed: false, + }, + offsets: SparseOffsetsCache::new(), + key_buffer: Vec::new(), + rows: Vec::new(), + last: None, + ordinal: 0, + leaves: Vec::new(), + leaf_bytes: 12, + descriptor_bytes: 0, + position: 0, + part_rows: 0, + part_blocks: 0, + manifest_bytes, + failed: false, + }) + } + async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { + require(batch.num_columns() >= 4, "samples require flat batches")?; + let column = |name: &str| { + batch + .column_by_name(name) + .ok_or_else(|| invalid(format!("missing sample input column {name}"))) + }; + let time = column(&self.manifest.timestamp.name)?; + require( + time.data_type() == &DataType::Timestamp(self.manifest.timestamp.unit.arrow(), None) + && time.null_count() == 0, + "samples timestamp type/null mismatch", + )?; + let (times, _) = timestamp_array_to_primitive(time) + .ok_or_else(|| invalid("invalid sample timestamp array"))?; + let sequence = column(consts::SEQUENCE_COLUMN_NAME)? + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid sample sequence type"))?; + let operations = column(consts::OP_TYPE_COLUMN_NAME)? + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid sample operation type"))?; + require( + sequence.null_count() == 0 && operations.null_count() == 0, + "null sample sequence/operation", + )?; + let fields = self + .manifest + .fields + .iter() + .map(|field| { + column(&field.name)? + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("invalid sample field array")) + }) + .collect::>>()?; + let keys = column(consts::PRIMARY_KEY_COLUMN_NAME)?; + require(keys.null_count() == 0, "null sample primary key")?; + let binary = keys.as_any().downcast_ref::(); + let dictionary = keys.as_any().downcast_ref::>(); + let values = dictionary.and_then(|d| d.values().as_any().downcast_ref::()); + require( + binary.is_some() || values.is_some(), + "samples require binary or UInt32 dictionary primary keys", + )?; + for i in 0..batch.num_rows() { + let key = if let Some(binary) = binary { + binary.value(i) + } else { + let dictionary = dictionary.ok_or_else(|| invalid("missing sample dictionary"))?; + let values = values.ok_or_else(|| invalid("invalid sample dictionary values"))?; + let index = dictionary.keys().value(i) as usize; + require( + index < values.len() && !values.is_null(index), + "invalid sample dictionary key", + )?; + 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 { + self.flush_block().await?; + self.ordinal = 0; + } + } + if self.rows.len() == self.options.block_rows + || !valid_rows(self.rows.len() + 1, fields.len()) + { + self.flush_block().await?; + } + self.rows.push(Row { + id, + time, + sequence: seq, + operation: op, + fields: fields + .iter() + .map(|field| { + if field.is_null(i) { + None + } else { + Some(field.value(i).to_bits()) + } + }) + .collect(), + }); + self.last = Some((id, time, seq)); + } + Ok(()) + } + async fn flush_block(&mut self) -> Result<()> { + if self.rows.is_empty() { + return Ok(()); + } + let rows = std::mem::take(&mut self.rows); + let (mut block, chunks) = encode_block(&rows, &self.manifest, self.ordinal)?; + self.ordinal = self + .ordinal + .checked_add(1) + .ok_or_else(|| invalid("sample ordinal overflow"))?; + let size = block.size(); + require( + size + 12 <= self.options.page_bytes, + "sample descriptor exceeds leaf limit", + )?; + if !self.leaves.is_empty() + && self.leaf_bytes + size > self.options.page_bytes + && self.leaves.len() == self.options.root_capacity + { + self.finish_part().await?; + } + if self.outputs.active.is_none() { + let id = u32::try_from(self.manifest.files.len()) + .map_err(|_| invalid("too many sample parts"))?; + self.outputs.open(self.manifest.part_path(id)).await?; + self.outputs.write(metadata::header()).await?; + self.position = metadata::HEADER_SIZE; + } + for (column, chunk) in block.columns.iter_mut().zip(chunks) { + column.payload.offset = self.position; + self.position = column.payload.end()?; + self.outputs.write(chunk).await?; + } + if self.leaves.is_empty() || self.leaf_bytes + size > self.options.page_bytes { + self.leaves.push(Vec::new()); + self.leaf_bytes = 12; + } + self.leaf_bytes += size; + self.descriptor_bytes += size; + self.part_rows = self + .part_rows + .checked_add(u64::from(block.rows)) + .ok_or_else(|| invalid("sample row count overflow"))?; + self.part_blocks += 1; + self.leaves + .last_mut() + .ok_or_else(|| invalid("missing sample leaf"))? + .push(block); + if self.position >= self.options.part_bytes + || self.descriptor_bytes >= self.options.metadata_bytes + { + self.finish_part().await?; + } + Ok(()) + } + async fn write_page(&mut self, page: Page) -> Result { + let bounds = page.bounds()?; + let (payload, bytes) = pack(&page.encode(), self.position, 1, true)?; + self.position = payload.end()?; + self.outputs.write(bytes).await?; + Ok(PageReference { bounds, payload }) + } + async fn finish_part(&mut self) -> Result<()> { + if self.leaves.is_empty() { + return Ok(()); + } + let payload_end = self.position; + let mut children = Vec::new(); + for leaf in std::mem::take(&mut self.leaves) { + children.push(self.write_page(Page::Leaf(leaf)).await?); + } + let depth = if children.len() == 1 { 1 } else { 2 }; + let root = if depth == 1 { + children.remove(0) + } else { + self.write_page(Page::Internal(children)).await? + }; + let bounds = root.bounds; + let file_length = self + .position + .checked_add(metadata::TRAILER_SIZE) + .ok_or_else(|| invalid("sample file length overflow"))?; + self.outputs + .write( + Trailer { + root, + depth, + file_length, + payload_end, + } + .encode(), + ) + .await?; + self.outputs.close().await?; + let part = Part { + part_id: u32::try_from(self.manifest.files.len()) + .map_err(|_| invalid("too many sample parts"))?, + file_size: file_length, + num_rows: self.part_rows, + num_blocks: self.part_blocks, + first: bounds.first, + last: bounds.last, + min_timestamp: bounds.min, + max_timestamp: bounds.max, + }; + self.manifest_bytes += serde_json::to_vec(&part).context(SerdeJsonSnafu)?.len() + + usize::from(!self.manifest.files.is_empty()); + require( + self.manifest_bytes <= MANIFEST_LIMIT, + "samples manifest exceeds size limit", + )?; + self.manifest.files.push(part); + self.position = 0; + self.part_rows = 0; + self.part_blocks = 0; + self.descriptor_bytes = 0; + self.leaf_bytes = 12; + Ok(()) + } + async fn complete(&mut self) -> Result { + require(!self.failed, "sample writer failed previously")?; + self.flush_block().await?; + self.finish_part().await?; + let bytes = serde_json::to_vec(&self.manifest).context(SerdeJsonSnafu)?; + require( + bytes.len() <= MANIFEST_LIMIT, + "samples manifest exceeds limit", + )?; + let mut file_size = bytes.len() as u64; + let mut objects = vec![ArtifactObject { + path: self.manifest.path(), + file_size, + }]; + for part in &self.manifest.files { + objects.push(ArtifactObject { + path: self.manifest.part_path(part.part_id), + file_size: part.file_size, + }); + file_size = file_size + .checked_add(part.file_size) + .ok_or_else(|| invalid("sample artifact size overflow"))?; + } + self.outputs.open(self.manifest.path()).await?; + self.outputs.write(bytes).await?; + self.outputs.close().await?; + self.outputs.committed = true; + Ok(ArtifactMetadata { + format: ArtifactFormat::ColumnarSamples, + generation: self.manifest.generation, + objects, + file_size, + }) + } +} +#[async_trait] +impl FormatWriter for SamplesWriter { + async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + require(!self.failed, "sample writer failed previously")?; + self.failed = true; + let result = self.write_batch(batch).await; + if result.is_ok() { + self.failed = false; + } + result + } + async fn finish(mut self: Box) -> Result { + let result = self.complete().await; + if result.is_err() + && let Err(error) = self.outputs.cleanup().await + { + common_telemetry::warn!(error; "Cleaning failed samples build"); + } + result + } + async fn abort(mut self: Box) -> Result<()> { + self.outputs.cleanup().await + } +} +fn encode_block(rows: &[Row], manifest: &Manifest, ordinal: u64) -> Result<(Block, 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, + min: first.time, + max: last.time, + ordinal, + rows: rows.len() as u32, + columns: Vec::new(), + }; + let mut chunks = Vec::new(); + let mut push = |role, id, encoding, bytes: Vec, nulls| -> Result<()> { + let (payload, bytes) = pack(&bytes, 0, encoding, false)?; + block.columns.push(Column { + role, + id, + payload, + nulls, + }); + chunks.push(bytes); + Ok(()) + }; + let (encoding, bytes) = codec::encode_time(&rows.iter().map(|r| r.time).collect::>()); + push(1, manifest.timestamp.column_id, encoding, bytes, 0)?; + let (encoding, bytes) = + codec::encode_sequence(&rows.iter().map(|r| r.sequence).collect::>()); + push(2, 0, encoding, bytes, 0)?; + let (encoding, bytes) = + codec::encode_operations(&rows.iter().map(|r| r.operation).collect::>()); + push(3, 0, encoding, bytes, 0)?; + for (index, field) in manifest.fields.iter().enumerate() { + let values: Vec<_> = rows.iter().map(|r| r.fields[index]).collect(); + let nulls = values.iter().filter(|v| v.is_none()).count(); + if nulls != rows.len() { + push( + 4, + field.column_id, + 5, + codec::encode_float(&values), + nulls as u32, + )?; + } + } + Ok((block, chunks)) +} diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index 145c2088cce..7a91295b739 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -67,11 +67,25 @@ pub(crate) async fn prepare_region_with_config( env: &mut TestEnv, timestamps: &[i64], config: MitoConfig, +) -> (MitoEngine, MitoRegionRef) { + prepare_region_with_metadata( + env, + timestamps, + config, + Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )), + ) + .await +} + +pub(crate) async fn prepare_region_with_metadata( + env: &mut TestEnv, + timestamps: &[i64], + config: MitoConfig, + metadata: store_api::metadata::RegionMetadataRef, ) -> (MitoEngine, MitoRegionRef) { let engine = env.create_engine(config).await; - let metadata = Arc::new(sst_region_metadata_with_encoding( - PrimaryKeyEncoding::Sparse, - )); let region_id = RegionId::new(1, 1); let mut request = CreateRequestBuilder::new().build(); request.column_metadatas = metadata.column_metadatas.clone(); @@ -129,7 +143,18 @@ pub(crate) async fn prepare_region_with_config( 0, )), ValueData::TimestampMillisecondValue(ts), - ValueData::U64Value(1), + if metadata + .field_columns() + .next() + .unwrap() + .column_schema + .data_type + == datatypes::prelude::ConcreteDataType::float64_datatype() + { + ValueData::F64Value(1.0) + } else { + ValueData::U64Value(1) + }, ])], }, hint: Some(WriteHint { diff --git a/src/promql/src/map.rs b/src/promql/src/map.rs index 0e2ff67fe02..a292234f350 100644 --- a/src/promql/src/map.rs +++ b/src/promql/src/map.rs @@ -43,6 +43,7 @@ use datatypes::series::time::{GridRange, TimeAxis}; use crate::functions::Round; use crate::grid_series::GridSeriesBuilder; +use crate::map_time::DateFunction; /// Float math supported by the evaluated-series kernel. #[derive(Clone, Copy, Debug)] @@ -88,6 +89,8 @@ pub enum MathOp { ClampMax { max: f64, }, + /// A date function of the value read as Unix seconds, such as `hour(v)`. + Date(DateFunction), } /// Float arithmetic between a vector and a scalar, as in Prometheus. @@ -301,6 +304,7 @@ impl MathOp { Self::Floor => v.floor(), Self::Sqrt => v.sqrt(), Self::Round { nearest } => Round::round_to_nearest(v, nearest), + Self::Date(function) => function.of_seconds(v), } } diff --git a/src/promql/src/map_time.rs b/src/promql/src/map_time.rs index c962f46be49..7177f436a1a 100644 --- a/src/promql/src/map_time.rs +++ b/src/promql/src/map_time.rs @@ -15,17 +15,13 @@ //! `time()` and the date functions over evaluated series, following Prometheus `funcTime` and //! `dateWrapper`. -use std::sync::Arc; - -use datafusion::arrow::array::{Array, Float64Array}; +use datafusion::arrow::array::Float64Array; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; -use datatypes::series::samples::{ - DenseSamplesRef, SampleRef, Samples, Series, SeriesFields, SeriesSchema, -}; +use datatypes::series::samples::{Samples, Series, SeriesFields, SeriesSchema}; use datatypes::series::time::{GridRange, TimeAxis}; -use crate::grid_series::{GridSeriesBuilder, METADATA_LABELS, external}; +use crate::grid_series::external; #[cfg(test)] mod tests; @@ -86,86 +82,6 @@ pub fn date_of_eval_time(function: DateFunction, schema: &SeriesSchema) -> Resul eval_time_series(schema, |ms| function.apply(ms / 1000)) } -/// A date function over `series`: each float value is read as Unix seconds and replaced by the -/// date component, and the metric name is dropped. Missing values stay missing; histogram -/// samples are dropped, so the result is `None` when no float remains. -/// -/// A missing value is a hole like any other, but a hole only skips a position: the result is one -/// series over the runs of consecutive positions that hold a value. -/// -/// Series that share a label set once their names are dropped must then go through -/// [`merge_same_labelset`](crate::grid_series::merge_same_labelset). -pub fn map_date( - function: DateFunction, - schema: &Arc, - series: &Series, -) -> Result> { - if series.time.grid().is_none() { - return Err(DataFusionError::Execution( - "date functions require an evaluated series".into(), - )); - } - let mut metric = series.metric.clone(); - metric.remove_all(&METADATA_LABELS); - let map = |seconds: f64| function.apply(go_int64(seconds)); - - // The fast path maps every float in place, so it needs each field to be entirely present: it - // keeps the input axis as it is. A hole means the mapped rows are fewer than the axis rows, - // so it falls through to the builder, which rebuilds the axis from the rows that survive. - let all_floats = series.fields.iter().all(|samples| { - samples.segments().iter().all(|segment| { - matches!(segment.dense(), Some(DenseSamplesRef::Float(floats)) if floats.null_count() == 0) - }) - }); - if all_floats { - let fields = series - .fields - .iter() - .map(|samples| { - let segments = samples - .segments() - .iter() - .map(|segment| { - let Some(DenseSamplesRef::Float(floats)) = segment.dense() else { - unreachable!("checked above"); - }; - Samples::from_float(floats.unary(map)) - }) - .collect(); - Samples::try_from_chunks(segments).map_err(external) - }) - .collect::>>()?; - let fields = SeriesFields::try_from_vec(fields).map_err(external)?; - let series = - Series::try_new(schema, metric, series.time.clone(), fields).map_err(external)?; - return Ok(Some(series)); - } - - // Holes in any field drop rows, so the builder names the positions that survive. - let mut builder = GridSeriesBuilder::try_new(schema.clone())?; - let mut values = vec![None; series.fields.len()]; - // The axis names the positions to evaluate and a cursor per field reads them in order, so an - // axis of several runs costs its rows and not a scan per row. - let mut cursors = series - .fields - .iter() - .map(|samples| samples.cursor()) - .collect::>(); - for (position, row) in series.time.positions() { - for (value, samples) in values.iter_mut().zip(&mut cursors) { - *value = match samples - .at(row) - .expect("the axis and the fields are aligned") - { - SampleRef::Float(seconds) => Some(map(seconds)), - SampleRef::Missing | SampleRef::Histogram(_) => None, - }; - } - builder.push(position, &values)?; - } - builder.finish(metric) -} - fn eval_time_series(schema: &SeriesSchema, value: impl Fn(i64) -> f64) -> Result { let grid = schema.eval_grid.clone().ok_or_else(|| { DataFusionError::Execution("time functions require an evaluation grid".into()) @@ -204,6 +120,25 @@ const UNIX_TO_ABSOLUTE: i64 = (1969 * 365 + 1969 / 4 - 1969 / 100 + 1969 / 400) + (ABSOLUTE_YEARS as i64 / 400 * 146_097 + MARCH_THRU_DECEMBER as i64) * 86_400; impl DateFunction { + /// The PromQL name of the function. + pub(crate) fn name(self) -> &'static str { + match self { + Self::DayOfMonth => "day_of_month", + Self::DayOfWeek => "day_of_week", + Self::DayOfYear => "day_of_year", + Self::DaysInMonth => "days_in_month", + Self::Hour => "hour", + Self::Minute => "minute", + Self::Month => "month", + Self::Year => "year", + } + } + + /// The component of `seconds` read as a Unix time, as Prometheus truncates it with `int64`. + pub(crate) fn of_seconds(self, seconds: f64) -> f64 { + self.apply(go_int64(seconds)) + } + /// The component of Unix second `seconds` in UTC, computed the way Go's `time` package does, /// including its wrapping arithmetic for instants out of its range. fn apply(self, seconds: i64) -> f64 { diff --git a/src/promql/src/map_time/tests.rs b/src/promql/src/map_time/tests.rs index 0124b420129..d90844d52b6 100644 --- a/src/promql/src/map_time/tests.rs +++ b/src/promql/src/map_time/tests.rs @@ -24,12 +24,13 @@ use datatypes::series::samples::{ }; use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis}; -use super::{DateFunction, QueryContextFunction, date_of_eval_time, map_date, time}; +use super::{DateFunction, QueryContextFunction, date_of_eval_time, time}; use crate::grid_series::merge_same_labelset; use crate::grid_series::reference::{ Labels, Points, canonical_bits, golden_cases, labels_of, metric_of, points_by_labels, points_of, run_ranges, series_of, sparse_series, }; +use crate::map::{MathOp, map_math}; fn schema(grid: Arc) -> Arc { Arc::new(SeriesSchema { @@ -71,7 +72,8 @@ fn time_functions_match_prometheus() { .iter() .filter_map(|(labels, points)| series_of(&schema, labels, points)) { - actual.extend(map_date(function(name), &schema, &series).unwrap()); + actual + .extend(map_math(&schema, &series, MathOp::Date(function(name))).unwrap()); } actual } @@ -107,7 +109,7 @@ fn time_functions_match_prometheus() { /// Histogram samples are dropped and the series trimmed to the remaining floats; chunked floats /// are mapped chunk by chunk. #[test] -fn map_date_drops_histograms_and_maps_chunks() { +fn date_functions_drop_histograms_and_map_chunks() { let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1000, 4).unwrap()).unwrap()); let schema = schema(grid.clone()); let range = GridRange::try_new(grid, 1, 3).unwrap(); @@ -155,7 +157,7 @@ fn map_date_drops_histograms_and_maps_chunks() { Samples::try_from_histogram(UInt32Array::from(vec![0]), histograms.finish().unwrap()) .unwrap(); let mixed = Samples::try_from_chunks(vec![histogram, floats(&[3600.0, 7200.0])]).unwrap(); - let hours = map_date(DateFunction::Hour, &schema, &series(mixed)) + let hours = map_math(&schema, &series(mixed), MathOp::Date(DateFunction::Hour)) .unwrap() .expect("a mixed series with one float"); assert_eq!( @@ -168,9 +170,13 @@ fn map_date_drops_histograms_and_maps_chunks() { ); let chunked = Samples::try_from_chunks(vec![floats(&[60.0]), floats(&[120.0, 180.0])]).unwrap(); - let minutes = map_date(DateFunction::Minute, &schema, &series(chunked)) - .unwrap() - .expect("a chunked float series"); + let minutes = map_math( + &schema, + &series(chunked), + MathOp::Date(DateFunction::Minute), + ) + .unwrap() + .expect("a chunked float series"); assert_eq!( points_of(&minutes), BTreeMap::from([ @@ -181,13 +187,13 @@ fn map_date_drops_histograms_and_maps_chunks() { ); } -/// `map_date` reads a series of several runs by the positions its axis carries. +/// A date function reads a series of several runs by the positions its axis carries. /// /// A sparse input cannot be mapped through the in-place fast path when a field holds a hole, so /// the piece-cutting path runs: the result must report the right positions, not an offset from /// the first one. #[test] -fn map_date_reads_a_sparse_axis_by_its_positions() { +fn date_functions_read_a_sparse_axis_by_its_positions() { let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1000, 8).unwrap()).unwrap()); let schema = schema(grid.clone()); // Positions 1, 2 and 6 hold seconds; the axis skips 3, 4 and 5. @@ -205,7 +211,7 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { assert_eq!(input.time.first_position(), Some(1)); assert_eq!(input.time.last_position(), Some(6)); - let output = map_date(DateFunction::Minute, &schema, &input) + let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute)) .unwrap() .expect("three rows"); // No field holds a hole, so the fast path keeps the axis as it is: one series, three rows. @@ -227,8 +233,8 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { // `minute()` of Unix seconds 1000, 2000 and 6000. assert_eq!(read, vec![(1, 16.0), (2, 33.0), (6, 40.0)]); - // A hole in the field keeps the fast path from applying, so the axis is rebuilt from the rows - // that survive: still one series, now over the runs of consecutive valued positions. + // A hole in a float field stays a missing row on the input's axis, as with every map, and the + // positions around it are mapped. let mut holed = Points::new(); for (position, seconds) in [(2, 120), (3, 180), (4, 240), (6, 360), (7, 420)] { holed.insert(position, f64::to_bits(seconds as f64)); @@ -239,10 +245,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { &holed, ) .unwrap(); - let output = map_date(DateFunction::Minute, &schema, &input) + let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute)) .unwrap() .expect("five rows"); - assert_eq!(run_ranges(&output), vec![(2, 3), (6, 2)]); + assert_eq!(run_ranges(&output), run_ranges(&input)); assert_eq!( points_of(&output), BTreeMap::from([ @@ -254,10 +260,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { ]) ); - // A sparse axis *and* a hole in the field, with the samples split over two chunks: the fast - // path cannot apply, so the piece-cutting path must walk the axis' positions and rebuild the - // axis from the rows that survive. Positions 0, 1, 3, 4 and 5 hold rows and position 3 is the - // hole, so a walk that counted rows from the first position would map the wrong seconds. + // A sparse axis *and* a hole in the field, with the samples split over two chunks: each chunk + // is mapped on the input's axis, the hole stays a missing row, and the values land on the + // positions the axis names. Positions 0, 1, 3, 4 and 5 hold rows and position 3 is the hole, + // so a walk that counted rows from the first position would map the wrong seconds. let time = TimeAxis::from_grid_spans(grid, vec![GridSpan::new(0, 2), GridSpan::new(3, 3)]).unwrap(); let chunked = Samples::try_from_chunks(vec![ @@ -273,11 +279,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { ) .unwrap(); assert_eq!(input.time.len(), 5); - let output = map_date(DateFunction::Minute, &schema, &input) + let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute)) .unwrap() - .expect("four rows"); - assert_eq!(output.time.len(), 4); - assert_eq!(run_ranges(&output), vec![(0, 2), (4, 2)]); + .expect("five rows"); + assert_eq!(run_ranges(&output), run_ranges(&input)); // `minute()` of Unix seconds 60, 120, 240 and 300. assert_eq!( points_of(&output), diff --git a/src/promql/src/native/display.rs b/src/promql/src/native/display.rs index c5db93d439e..28099e6c50b 100644 --- a/src/promql/src/native/display.rs +++ b/src/promql/src/native/display.rs @@ -272,6 +272,7 @@ fn math_op(op: &MathOp) -> String { MathOp::Clamp { min, max } => format!("clamp {min} {max}"), MathOp::ClampMin { min } => format!("clamp_min {min}"), MathOp::ClampMax { max } => format!("clamp_max {max}"), + MathOp::Date(function) => function.name().to_string(), } } diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index d388faed5ac..cc79b2b8fa2 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -4158,9 +4158,10 @@ impl NativeExecutor { /// Maps one chunk with `ops`, keeping its axis and its closing. /// -/// Mapping is point-wise over valued positions, so a chunk whose values survive keeps exactly the -/// rows it came with: a position no field fills is not a row of any chunk, and mapping does not -/// invent one. `None` means the operation emptied every position, which yields no chunk at all +/// Mapping is point-wise over valued positions: a float chunk keeps the rows it came with, and a +/// chunk holding native histograms loses the rows only histograms filled, as math skips them; a +/// position no field fills is not a row of any chunk, and mapping does not invent one. `None` means +/// the operation emptied every position, which yields no chunk at all /// unless this chunk closed its series: a closing chunk is still handed over, holding no position, /// because a consumer tells "nothing more here" from "no more input" by that chunk alone. fn map_chunk( @@ -4172,7 +4173,8 @@ fn map_chunk( for &op in ops { series = match map_math(schema, &series, op)? { Some(mapped) => mapped, - // `clamp` with `max < min` is the empty vector: every position it held was decided and + // `clamp` with `max < min` is the empty vector, and math skips native histograms, so a + // series without a float value has none left: every position it held was decided and // holds no value now, so there is no row to hand over. None => { return if chunk.is_last { diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 83f3e95e98a..53bb3a187a0 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -127,6 +127,7 @@ use datatypes::series::time::{EvalGrid, GridSpec}; use datatypes::series::{MetricName, SeriesOrdering}; use promql::histogram_eval::HistogramFunction; use promql::map::{ArithmeticOp, ComparisonOp, MathOp}; +use promql::map_time::DateFunction; use promql::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, @@ -1310,6 +1311,18 @@ fn math_signature(name: &str) -> Option { "deg" => Some(MathSignature::Unary(MathOp::Deg)), "rad" => Some(MathSignature::Unary(MathOp::Rad)), "sgn" => Some(MathSignature::Unary(MathOp::Sgn)), + // A date function takes the vector its values are read from; without one it reads the + // evaluation time, which has no input to map and is refused by `math_op`. + "day_of_month" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfMonth))), + "day_of_week" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfWeek))), + "day_of_year" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfYear))), + "days_in_month" => Some(MathSignature::Unary(MathOp::Date( + DateFunction::DaysInMonth, + ))), + "hour" => Some(MathSignature::Unary(MathOp::Date(DateFunction::Hour))), + "minute" => Some(MathSignature::Unary(MathOp::Date(DateFunction::Minute))), + "month" => Some(MathSignature::Unary(MathOp::Date(DateFunction::Month))), + "year" => Some(MathSignature::Unary(MathOp::Date(DateFunction::Year))), // `round`'s nearest defaults to one, so `round(x)` is `round(x, 1)` and its variant always // carries the number it rounds to. "round" => Some(MathSignature::DefaultedConstant { @@ -1347,6 +1360,12 @@ fn math_op<'a>( // The vector is the function's first argument, so a call that wrote no argument at all is // refused by the same arity check as one that wrote too few constants. let Some((vector, parameters)) = arguments.split_first() else { + if let MathSignature::Unary(MathOp::Date(_)) = signature { + return Err(not_implemented(&format!( + "a call to `{name}` without an argument, which reads the evaluation time and has \ + no native kernel" + ))); + } return Err(arity_error(name, signature, 0)); }; let given = parameters.len() + 1; diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 38bb8a6854b..d131af69ba9 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -537,6 +537,9 @@ fn a_call_with_the_wrong_number_of_arguments_names_the_count_it_takes() { "abs", "with 0 argument(s), where it takes its vector", ), + // A date function written without an argument is valid PromQL that reads the evaluation + // time; there is no input to map, so it is left to the ordinary path. + (eval_stmt("hour()"), "hour", "reads the evaluation time"), ]; for (stmt, name, expected) in cases { @@ -1212,6 +1215,15 @@ fn the_supported_math_functions_map_to_their_kernels() { ("clamp(x, 0, 1)", "clamp 0 1"), ("clamp_min(x, 0)", "clamp_min 0"), ("clamp_max(x, 1)", "clamp_max 1"), + // Date functions over a vector read its values as Unix seconds. + ("day_of_month(x)", "day_of_month"), + ("day_of_week(x)", "day_of_week"), + ("day_of_year(x)", "day_of_year"), + ("days_in_month(x)", "days_in_month"), + ("hour(x)", "hour"), + ("minute(x)", "minute"), + ("month(x)", "month"), + ("year(x)", "year"), ]; for (query, op) in cases { diff --git a/tests/compatibility/wire/mito_columnar_samples/README.md b/tests/compatibility/wire/mito_columnar_samples/README.md new file mode 100644 index 00000000000..3499a6fce02 --- /dev/null +++ b/tests/compatibility/wire/mito_columnar_samples/README.md @@ -0,0 +1,39 @@ +# Mito columnar samples version 1 + +This fixture freezes the version-1 sample wire contract in +[`mito-samples.md`](../../../../docs/formats/mito-samples.md). It is exercised by +the owning Mito Rust test, not SQLness discovery: + +```shell +cargo nextest run -p mito2 frozen_v1_samples_fixture_and_unknown_codec +``` + +`v1-manifest.json` and `v1-part.hex` describe a 397-byte part with raw compression, +one leaf/root, and two samples for `(table_id=7, tsid=42)` at timestamp 100 ms. +The first row is a put with sequence 9 and Float64 negative zero. The second is a +delete with sequence 8 and a null field. Both rows must survive decoding. + +The part was assembled independently of the Rust writer, using explicit +little-endian fields and XXH64 with seed zero. Its layout is: + +| Range | Content | +| --- | --- | +| 0..16 | Header | +| 16..32 | Constant-step timestamps: first 100, step 0 | +| 32..41 | Delta sequences: first 9, ZigZag delta -1 | +| 41..43 | Raw operations: put 1, delete 0 | +| 43..53 | Nullable XOR: bitmap mode, validity 0x01, negative-zero bits | +| 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. + +Do not regenerate this fixture just to make a reader change pass. Preserve its +bytes for future reader compatibility; introduce a separate versioned fixture +when the approved format changes. This is a frozen-byte compatibility check, +not a cross-binary upgrade/downgrade test. Older binaries are not promised to +understand the new opt-in `columnar-samples` catalog identifier. diff --git a/tests/compatibility/wire/mito_columnar_samples/v1-manifest.json b/tests/compatibility/wire/mito_columnar_samples/v1-manifest.json new file mode 100644 index 00000000000..0a62f7929d0 --- /dev/null +++ b/tests/compatibility/wire/mito_columnar_samples/v1-manifest.json @@ -0,0 +1,38 @@ +{ + "format": "columnar-samples", + "version": 1, + "region_id": 4294967297, + "source_file_id": "11111111-1111-1111-1111-111111111111", + "generation": "22222222-2222-2222-2222-222222222222", + "timestamp": { + "column_id": 3, + "name": "ts", + "unit": "millisecond" + }, + "fields": [ + { + "column_id": 2, + "name": "field_0", + "type": "float64", + "nullable": true + } + ], + "files": [ + { + "part_id": 0, + "file_size": 397, + "num_rows": 2, + "num_blocks": 1, + "first": { + "table_id": 7, + "tsid": 42 + }, + "last": { + "table_id": 7, + "tsid": 42 + }, + "min_timestamp": 100, + "max_timestamp": 100 + } + ] +} diff --git a/tests/compatibility/wire/mito_columnar_samples/v1-part.hex b/tests/compatibility/wire/mito_columnar_samples/v1-part.hex new file mode 100644 index 00000000000..8e27c6d801d --- /dev/null +++ b/tests/compatibility/wire/mito_columnar_samples/v1-part.hex @@ -0,0 +1,13 @@ +475453414d504c45010000000000000064000000000000000000000000000000 +0900000000000000010100010100000000000000800000000001000000000000 +00070000002a0000000000000064000000000000006400000000000000000000 +0000000000020000000400000001000000030000001000000000000000100000 +00100000005ec3a3cc10ea641301000000000000000000000002000000000000 +00200000000000000009000000090000003b1e9c01038b57ca04000000000000 +0000000000030000000000000029000000000000000200000002000000f1577f +bf1d628fdd07000000000000000000000004000000020000002b000000000000 +000a0000000a0000008e2dc4e1ef3f8b09050000000000000001000000070000 +002a00000000000000070000002a000000000000006400000000000000640000 +00000000003500000000000000e8000000e800000077c4ce1b4e0f6c87010000 +000000000001000000010000008d010000000000003500000000000000475453 +414d504c45f202e0cce4a9fdb1